From 36dfff1743c238857719341bb187965f8beab2fe Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Wed, 16 Sep 2026 07:37:44 +0530 Subject: [PATCH] refactor(tool): migrate Batch/Cron/NotebookEdit/Config/Impact/GitHistory to typed schemas Follow the Schema()+Parameters() single-source pattern with wire-shape parity tests (including nested object and enum preservation). Execute parses via DecodeInput[T]. --- internal/tool/batch.go | 50 ++++++++------ internal/tool/cron.go | 106 +++++++++++++++-------------- internal/tool/extra.go | 80 ++++++++++++++-------- internal/tool/git_history.go | 57 ++++++++-------- internal/tool/impact.go | 48 ++++++------- internal/tool/schema_batch_test.go | 104 ++++++++++++++++++++++++++++ internal/tool/tool.go | 16 +++++ 7 files changed, 311 insertions(+), 150 deletions(-) diff --git a/internal/tool/batch.go b/internal/tool/batch.go index 3799b206b..8af59fe08 100644 --- a/internal/tool/batch.go +++ b/internal/tool/batch.go @@ -26,38 +26,50 @@ func (BatchTool) Description() string { return "Run several read-only tool calls in one turn (fan-out research). Calls must be read-only tools." } -func (BatchTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "calls": map[string]interface{}{ - "type": "array", - "description": "Read-only tool calls to execute in sequence", - "items": map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "tool": map[string]interface{}{"type": "string", "description": "Read-only tool name (e.g. Read, Grep, Glob, LS, CodeSearch)"}, - "input": map[string]interface{}{"type": "object", "description": "Tool input object"}, +// BatchInput is the typed input for BatchTool. +type BatchInput struct { + Calls []batchCall `json:"calls"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (BatchTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "calls": { + Type: "array", + Description: "Read-only tool calls to execute in sequence", + Items: &SchemaProperty{ + Type: "object", + Properties: map[string]SchemaProperty{ + "tool": {Type: "string", Description: "Read-only tool name (e.g. Read, Grep, Glob, LS, CodeSearch)"}, + "input": {Type: "object", Description: "Tool input object"}, }, - "required": []string{"tool"}, + Required: []string{"tool"}, }, }, }, - "required": []string{"calls"}, + Required: []string{"calls"}, } } +func (BatchTool) Parameters() map[string]interface{} { + return batchSchema.ToJSONSchema() +} + +// batchSchema is the single source of truth for Batch's input schema. +var batchSchema = BatchTool{}.Schema() + type batchCall struct { Tool string `json:"tool"` Input json.RawMessage `json:"input"` } func (BatchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Calls []batchCall `json:"calls"` - } - if err := json.Unmarshal(input, &p); err != nil { - return "", fmt.Errorf("batch: %w", err) + p, err := DecodeInput[BatchInput]("Batch", input) + if err != nil { + return "", err } if len(p.Calls) == 0 { return "", errors.New("batch: calls is required and must not be empty") diff --git a/internal/tool/cron.go b/internal/tool/cron.go index 7fd0c428e..4d450a663 100644 --- a/internal/tool/cron.go +++ b/internal/tool/cron.go @@ -220,49 +220,43 @@ func (CronCreateTool) Description() string { return "Schedule a prompt to run at a future time — either recurring on a cron schedule, or once at a specific time." } -func (CronCreateTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "schedule": map[string]interface{}{ - "type": "string", - "description": "5-field cron expression in user's local timezone: minute hour day-of-month month day-of-week", - }, - "prompt": map[string]interface{}{ - "type": "string", - "description": "The prompt to enqueue when the schedule fires", - }, - "recurring": map[string]interface{}{ - "type": "boolean", - "description": "If true, repeats on schedule. If false, fires once then auto-deletes (default: true)", - }, - "durable": map[string]interface{}{ - "type": "boolean", - "description": "If true, persists to disk and survives session restarts (default: false)", - }, - "max_runs": map[string]interface{}{ - "type": "integer", - "description": "Stop after this many fires (0 = unlimited). Useful for /loop-style caps.", - }, - "expires_in_sec": map[string]interface{}{ - "type": "integer", - "description": "Auto-delete job after this many seconds from creation (0 = never).", - }, +// CronCreateInput is the typed input for CronCreateTool. +type CronCreateInput struct { + Schedule string `json:"schedule"` + Prompt string `json:"prompt"` + Recurring *bool `json:"recurring"` + Durable *bool `json:"durable"` + MaxRuns int `json:"max_runs"` + ExpiresInSec int `json:"expires_in_sec"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (CronCreateTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "schedule": {Type: "string", Description: "5-field cron expression in user's local timezone: minute hour day-of-month month day-of-week"}, + "prompt": {Type: "string", Description: "The prompt to enqueue when the schedule fires"}, + "recurring": {Type: "boolean", Description: "If true, repeats on schedule. If false, fires once then auto-deletes (default: true)"}, + "durable": {Type: "boolean", Description: "If true, persists to disk and survives session restarts (default: false)"}, + "max_runs": {Type: "integer", Description: "Stop after this many fires (0 = unlimited). Useful for /loop-style caps."}, + "expires_in_sec": {Type: "integer", Description: "Auto-delete job after this many seconds from creation (0 = never)."}, }, - "required": []string{"schedule", "prompt"}, + Required: []string{"schedule", "prompt"}, } } +func (CronCreateTool) Parameters() map[string]interface{} { + return cronCreateSchema.ToJSONSchema() +} + +// cronCreateSchema is the single source of truth for CronCreate's input schema. +var cronCreateSchema = CronCreateTool{}.Schema() + func (CronCreateTool) Execute(_ context.Context, input json.RawMessage) (string, error) { - var p struct { - Schedule string `json:"schedule"` - Prompt string `json:"prompt"` - Recurring *bool `json:"recurring"` - Durable *bool `json:"durable"` - MaxRuns int `json:"max_runs"` - ExpiresInSec int `json:"expires_in_sec"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[CronCreateInput]("CronCreate", input) + if err != nil { return "", err } if p.Schedule == "" { @@ -317,24 +311,34 @@ type CronDeleteTool struct{} func (CronDeleteTool) Name() string { return "CronDelete" } func (CronDeleteTool) Aliases() []string { return []string{"cron_delete"} } func (CronDeleteTool) Description() string { return "Remove a scheduled cron job" } -func (CronDeleteTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "id": map[string]interface{}{ - "type": "string", - "description": "The cron job ID to delete", - }, + +// CronDeleteInput is the typed input for CronDeleteTool. +type CronDeleteInput struct { + ID string `json:"id"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (CronDeleteTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "id": {Type: "string", Description: "The cron job ID to delete"}, }, - "required": []string{"id"}, + Required: []string{"id"}, } } +func (CronDeleteTool) Parameters() map[string]interface{} { + return cronDeleteSchema.ToJSONSchema() +} + +// cronDeleteSchema is the single source of truth for CronDelete's input schema. +var cronDeleteSchema = CronDeleteTool{}.Schema() + func (CronDeleteTool) Execute(_ context.Context, input json.RawMessage) (string, error) { - var p struct { - ID string `json:"id"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[CronDeleteInput]("CronDelete", input) + if err != nil { return "", err } if !globalCronScheduler.Delete(p.ID) { diff --git a/internal/tool/extra.go b/internal/tool/extra.go index a3666132e..d0bc020d8 100644 --- a/internal/tool/extra.go +++ b/internal/tool/extra.go @@ -17,25 +17,37 @@ func (NotebookEditTool) Description() string { return "Edit a Jupyter notebook cell. Specify the notebook path, cell number, and new source." } -func (NotebookEditTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "path": map[string]interface{}{"type": "string", "description": "Notebook file path (.ipynb)"}, - "cell_number": map[string]interface{}{"type": "integer", "description": "Cell number (0-based)"}, - "new_source": map[string]interface{}{"type": "string", "description": "New cell source content"}, +// NotebookEditInput is the typed input for NotebookEditTool. +type NotebookEditInput struct { + Path string `json:"path"` + CellNumber int `json:"cell_number"` + NewSource string `json:"new_source"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (NotebookEditTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "path": {Type: "string", Description: "Notebook file path (.ipynb)"}, + "cell_number": {Type: "integer", Description: "Cell number (0-based)"}, + "new_source": {Type: "string", Description: "New cell source content"}, }, - "required": []string{"path", "cell_number", "new_source"}, + Required: []string{"path", "cell_number", "new_source"}, } } +func (NotebookEditTool) Parameters() map[string]interface{} { + return notebookEditSchema.ToJSONSchema() +} + +// notebookEditSchema is the single source of truth for NotebookEdit's input schema. +var notebookEditSchema = NotebookEditTool{}.Schema() + func (NotebookEditTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Path string `json:"path"` - CellNumber int `json:"cell_number"` - NewSource string `json:"new_source"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[NotebookEditInput]("NotebookEdit", input) + if err != nil { return "", err } if err := validatePathAllowed(ctx, p.Path); err != nil { @@ -92,25 +104,37 @@ func (ConfigTool) Description() string { return "Read or modify rho configuration settings." } -func (ConfigTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "action": map[string]interface{}{"type": "string", "enum": []string{"get", "set"}, "description": "Action"}, - "key": map[string]interface{}{"type": "string", "description": "Setting key"}, - "value": map[string]interface{}{"type": "string", "description": "Setting value (for set)"}, +// ConfigInput is the typed input for ConfigTool. +type ConfigInput struct { + Action string `json:"action"` + Key string `json:"key"` + Value string `json:"value"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (ConfigTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "action": {Type: "string", Description: "Action", Enum: []interface{}{"get", "set"}}, + "key": {Type: "string", Description: "Setting key"}, + "value": {Type: "string", Description: "Setting value (for set)"}, }, - "required": []string{"action", "key"}, + Required: []string{"action", "key"}, } } +func (ConfigTool) Parameters() map[string]interface{} { + return configSchema.ToJSONSchema() +} + +// configSchema is the single source of truth for Config's input schema. +var configSchema = ConfigTool{}.Schema() + func (ConfigTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Action string `json:"action"` - Key string `json:"key"` - Value string `json:"value"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[ConfigInput]("Config", input) + if err != nil { return "", err } tc := GetToolContext(ctx) diff --git a/internal/tool/git_history.go b/internal/tool/git_history.go index 8412cf649..68749fa71 100644 --- a/internal/tool/git_history.go +++ b/internal/tool/git_history.go @@ -21,40 +21,39 @@ func (GitHistoryTool) Description() string { return "Mine git history for co-change patterns, file history, and code ownership insights." } -func (GitHistoryTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "action": map[string]interface{}{ - "type": "string", - "enum": []string{"history", "cochange", "owners", "blame"}, - "description": "Action: history (commit log), cochange (frequently changed together), owners (top contributors), blame (line-level ownership)", - }, - "file": map[string]interface{}{ - "type": "string", - "description": "File to analyze", - }, - "limit": map[string]interface{}{ - "type": "integer", - "description": "Max results (default: 20)", - }, - "root": map[string]interface{}{ - "type": "string", - "description": "Project root directory (default: current dir)", - }, +// GitHistoryInput is the typed input for GitHistoryTool. +type GitHistoryInput struct { + Action string `json:"action"` + File string `json:"file"` + Limit int `json:"limit"` + Root string `json:"root"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (GitHistoryTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "action": {Type: "string", Description: "Action: history (commit log), cochange (frequently changed together), owners (top contributors), blame (line-level ownership)", Enum: []interface{}{"history", "cochange", "owners", "blame"}}, + "file": {Type: "string", Description: "File to analyze"}, + "limit": {Type: "integer", Description: "Max results (default: 20)"}, + "root": {Type: "string", Description: "Project root directory (default: current dir)"}, }, - "required": []string{"action"}, + Required: []string{"action"}, } } +func (GitHistoryTool) Parameters() map[string]interface{} { + return gitHistorySchema.ToJSONSchema() +} + +// gitHistorySchema is the single source of truth for GitHistory's input schema. +var gitHistorySchema = GitHistoryTool{}.Schema() + func (GitHistoryTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Action string `json:"action"` - File string `json:"file"` - Limit int `json:"limit"` - Root string `json:"root"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[GitHistoryInput]("GitHistory", input) + if err != nil { return "", err } diff --git a/internal/tool/impact.go b/internal/tool/impact.go index 5d7c246fa..de103320d 100644 --- a/internal/tool/impact.go +++ b/internal/tool/impact.go @@ -24,34 +24,36 @@ func (ImpactTool) Description() string { return "Analyze cross-file impact of changes: dependencies, dependents, co-change patterns, risk score, and test suggestions." } -func (ImpactTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "files": map[string]interface{}{ - "type": "array", - "items": map[string]interface{}{"type": "string"}, - "description": "Files to analyze impact for (e.g. changed files). If empty, uses git diff.", - }, - "depth": map[string]interface{}{ - "type": "integer", - "description": "Max traversal depth for dependency analysis (default: 3)", - }, - "root": map[string]interface{}{ - "type": "string", - "description": "Project root directory (default: current dir)", - }, +// ImpactInput is the typed input for ImpactTool. +type ImpactInput struct { + Files []string `json:"files"` + Depth int `json:"depth"` + Root string `json:"root"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (ImpactTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "files": {Type: "array", Description: "Files to analyze impact for (e.g. changed files). If empty, uses git diff.", Items: &SchemaProperty{Type: "string"}}, + "depth": {Type: "integer", Description: "Max traversal depth for dependency analysis (default: 3)"}, + "root": {Type: "string", Description: "Project root directory (default: current dir)"}, }, } } +func (ImpactTool) Parameters() map[string]interface{} { + return impactSchema.ToJSONSchema() +} + +// impactSchema is the single source of truth for Impact's input schema. +var impactSchema = ImpactTool{}.Schema() + func (ImpactTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Files []string `json:"files"` - Depth int `json:"depth"` - Root string `json:"root"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[ImpactInput]("Impact", input) + if err != nil { return "", err } diff --git a/internal/tool/schema_batch_test.go b/internal/tool/schema_batch_test.go index 12fed7cc0..10ad8f11d 100644 --- a/internal/tool/schema_batch_test.go +++ b/internal/tool/schema_batch_test.go @@ -175,3 +175,107 @@ func TestFileEditSchemaProvider(t *testing.T) { t.Fatalf("required = %v, want [path old_str]", FileEditTool{}.Parameters()["required"]) } } + +func TestBatchSchemaProvider(t *testing.T) { + var _ SchemaProvider = BatchTool{} + props := schemaProps(t, BatchTool{}.Parameters()) + calls, ok := props["calls"].(map[string]interface{}) + if !ok || calls["type"] != "array" { + t.Fatalf("calls prop = %v, want array type", props["calls"]) + } + items, ok := calls["items"].(map[string]interface{}) + if !ok || items["type"] != "object" { + t.Fatalf("calls items = %v, want object type", calls["items"]) + } + sub, ok := items["properties"].(map[string]interface{}) + if !ok { + t.Fatal("calls items properties missing") + } + if sub["tool"].(map[string]interface{})["type"] != "string" { + t.Fatal("calls items tool type wrong") + } + if req, _ := items["required"].([]string); len(req) != 1 || req[0] != "tool" { + t.Fatalf("calls items required = %v, want [tool]", items["required"]) + } + req, _ := BatchTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "calls" { + t.Fatalf("required = %v, want [calls]", BatchTool{}.Parameters()["required"]) + } +} + +func TestCronSchemasProvider(t *testing.T) { + var _ SchemaProvider = CronCreateTool{} + var _ SchemaProvider = CronDeleteTool{} + props := schemaProps(t, CronCreateTool{}.Parameters()) + if props["schedule"].(map[string]interface{})["type"] != "string" { + t.Fatal("schedule type wrong") + } + req, _ := CronCreateTool{}.Parameters()["required"].([]string) + if len(req) != 2 || req[0] != "schedule" || req[1] != "prompt" { + t.Fatalf("required = %v, want [schedule prompt]", CronCreateTool{}.Parameters()["required"]) + } + dprops := schemaProps(t, CronDeleteTool{}.Parameters()) + if dprops["id"].(map[string]interface{})["type"] != "string" { + t.Fatal("id type wrong") + } +} + +func TestNotebookEditSchemaProvider(t *testing.T) { + var _ SchemaProvider = NotebookEditTool{} + props := schemaProps(t, NotebookEditTool{}.Parameters()) + for _, f := range []string{"path", "new_source"} { + if props[f].(map[string]interface{})["type"] != "string" { + t.Fatalf("%s type wrong", f) + } + } + if props["cell_number"].(map[string]interface{})["type"] != "integer" { + t.Fatal("cell_number type wrong") + } +} + +func TestConfigSchemaProvider(t *testing.T) { + var _ SchemaProvider = ConfigTool{} + props := schemaProps(t, ConfigTool{}.Parameters()) + action, ok := props["action"].(map[string]interface{}) + if !ok || action["type"] != "string" { + t.Fatalf("action prop = %v, want string type", props["action"]) + } + enum, ok := action["enum"].([]interface{}) + if !ok || len(enum) != 2 { + t.Fatalf("action enum = %v, want [get set]", action["enum"]) + } +} + +func TestImpactSchemaProvider(t *testing.T) { + var _ SchemaProvider = ImpactTool{} + props := schemaProps(t, ImpactTool{}.Parameters()) + files, ok := props["files"].(map[string]interface{}) + if !ok || files["type"] != "array" { + t.Fatalf("files prop = %v, want array type", props["files"]) + } + items, ok := files["items"].(map[string]interface{}) + if !ok || items["type"] != "string" { + t.Fatalf("files items = %v, want string type", files["items"]) + } + _, hasRequired := ImpactTool{}.Parameters()["required"] + if hasRequired { + t.Fatalf("Impact has no required fields, got %v", ImpactTool{}.Parameters()["required"]) + } +} + +func TestGitHistorySchemaProvider(t *testing.T) { + var _ SchemaProvider = GitHistoryTool{} + props := schemaProps(t, GitHistoryTool{}.Parameters()) + action, ok := props["action"].(map[string]interface{}) + if !ok || action["type"] != "string" { + t.Fatalf("action prop = %v, want string type", props["action"]) + } + enum, ok := action["enum"].([]interface{}) + if !ok || len(enum) != 4 { + t.Fatalf("action enum = %v, want 4 options", action["enum"]) + } + req, _ := GitHistoryTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "action" { + t.Fatalf("required = %v, want [action]", GitHistoryTool{}.Parameters()["required"]) + } +} diff --git a/internal/tool/tool.go b/internal/tool/tool.go index 15b72e9b9..f2f423a64 100644 --- a/internal/tool/tool.go +++ b/internal/tool/tool.go @@ -58,6 +58,12 @@ type SchemaProperty struct { Enum []interface{} `json:"enum,omitempty"` // Items describes the element type for array fields. Items *SchemaProperty `json:"items,omitempty"` + // Properties describes the fields of a nested object (used together with + // Type "object", e.g. array items that are objects with their own + // required fields). + Properties map[string]SchemaProperty `json:"properties,omitempty"` + // Required lists required field names for a nested object. + Required []string `json:"required,omitempty"` // Minimum/Maximum constrain numeric fields. Kept as interface{} (like // Default) so authors write plain literals and the wire output matches // hand-written schemas exactly. @@ -111,6 +117,16 @@ func (p SchemaProperty) toMap() map[string]interface{} { if p.Items != nil { m["items"] = p.Items.toMap() } + if len(p.Properties) > 0 { + props := make(map[string]interface{}, len(p.Properties)) + for name, sub := range p.Properties { + props[name] = sub.toMap() + } + m["properties"] = props + } + if len(p.Required) > 0 { + m["required"] = p.Required + } return m }