From 9e2e411b68764199a485a69f9e88d7f30138fe4c Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Wed, 16 Sep 2026 18:03:43 +0530 Subject: [PATCH] refactor(tool): extend SchemaProperty and migrate WebSearch, MultiAgent to typed schemas --- internal/tool/agent.go | 71 +++++++++++++---------- internal/tool/schema_batch_test.go | 91 ++++++++++++++++++++++++++++++ internal/tool/tool.go | 44 ++++++++++++++- internal/tool/web_search.go | 67 ++++++++++------------ 4 files changed, 203 insertions(+), 70 deletions(-) diff --git a/internal/tool/agent.go b/internal/tool/agent.go index 45809d88..c12cda3b 100644 --- a/internal/tool/agent.go +++ b/internal/tool/agent.go @@ -210,47 +210,56 @@ func (MultiAgentTool) Description() string { "Each task is a prompt string (explore by default) or an object with typed spawn fields." } -func (MultiAgentTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "tasks": map[string]interface{}{ - "type": "array", - "items": map[string]interface{}{ - "oneOf": []interface{}{ - map[string]interface{}{"type": "string"}, - map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "prompt": map[string]interface{}{"type": "string"}, - "subagent_type": map[string]interface{}{"type": "string"}, - "capability_mode": map[string]interface{}{"type": "string"}, - "isolation": map[string]interface{}{"type": "string"}, - "thoroughness": map[string]interface{}{"type": "string"}, - "description": map[string]interface{}{"type": "string"}, - "model": map[string]interface{}{"type": "string"}, - "cwd": map[string]interface{}{"type": "string"}, +// MultiAgentInput is the typed input for MultiAgentTool. +type MultiAgentInput struct { + Tasks []json.RawMessage `json:"tasks"` + RunInBackground bool `json:"run_in_background"` +} + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (MultiAgentTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "tasks": { + Type: "array", + Items: &SchemaProperty{ + OneOf: []SchemaProperty{ + {Type: "string"}, + { + Type: "object", + Properties: map[string]SchemaProperty{ + "prompt": {Type: "string"}, + "subagent_type": {Type: "string"}, + "capability_mode": {Type: "string"}, + "isolation": {Type: "string"}, + "thoroughness": {Type: "string"}, + "description": {Type: "string"}, + "model": {Type: "string"}, + "cwd": {Type: "string"}, }, - "required": []string{"prompt"}, + Required: []string{"prompt"}, }, }, }, }, - "run_in_background": map[string]interface{}{ - "type": "boolean", - "description": "If true, spawn all sub-agents in the background.", - }, + "run_in_background": {Type: "boolean", Description: "If true, spawn all sub-agents in the background."}, }, - "required": []string{"tasks"}, + Required: []string{"tasks"}, } } +func (MultiAgentTool) Parameters() map[string]interface{} { + return multiAgentSchema.ToJSONSchema() +} + +// multiAgentSchema is the single source of truth for MultiAgent's input schema. +var multiAgentSchema = MultiAgentTool{}.Schema() + func (MultiAgentTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Tasks []json.RawMessage `json:"tasks"` - RunInBackground bool `json:"run_in_background"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[MultiAgentInput]("MultiAgent", input) + if err != nil { return "", err } if len(p.Tasks) > maxParallelAgentTasks { diff --git a/internal/tool/schema_batch_test.go b/internal/tool/schema_batch_test.go index 5125aef2..5b92f98a 100644 --- a/internal/tool/schema_batch_test.go +++ b/internal/tool/schema_batch_test.go @@ -1379,3 +1379,94 @@ func TestTerminalKillSchemaProvider(t *testing.T) { t.Fatalf("required = %v, want [terminal_id]", TerminalKillTool{}.Parameters()["required"]) } } + +func TestWebSearchSchemaProvider(t *testing.T) { + var _ SchemaProvider = WebSearchTool{} + props := schemaProps(t, WebSearchTool{}.Parameters()) + query := props["query"].(map[string]interface{}) + if query["type"] != "string" { + t.Fatal("query type wrong") + } + if query["maxLength"] != 2000 { + t.Fatalf("query maxLength = %v, want 2000", query["maxLength"]) + } + queries := props["queries"].(map[string]interface{}) + if queries["type"] != "array" { + t.Fatalf("queries type = %v, want array", queries["type"]) + } + if queries["maxItems"] != 20 { + t.Fatalf("queries maxItems = %v, want 20", queries["maxItems"]) + } + items := queries["items"].(map[string]interface{}) + if items["type"] != "string" || items["maxLength"] != 2000 { + t.Fatalf("queries items = %v, want string+maxLength 2000", items) + } + num := props["numResults"].(map[string]interface{}) + if num["default"] != 5 || num["minimum"] != 1 || num["maximum"] != 20 { + t.Fatalf("numResults = %v, want min/max/default", num) + } + st := props["searchType"].(map[string]interface{}) + enum, ok := st["enum"].([]interface{}) + if !ok || len(enum) != 2 || enum[0] != "web" || enum[1] != "news" { + t.Fatalf("searchType enum = %v, want [web news]", st["enum"]) + } + if st["default"] != "web" { + t.Fatalf("searchType default = %v, want web", st["default"]) + } +} + +func TestMultiAgentSchemaProvider(t *testing.T) { + var _ SchemaProvider = MultiAgentTool{} + props := schemaProps(t, MultiAgentTool{}.Parameters()) + tasks := props["tasks"].(map[string]interface{}) + if tasks["type"] != "array" { + t.Fatalf("tasks type = %v, want array", tasks["type"]) + } + if _, hasType := tasks["items"].(map[string]interface{})["type"]; hasType { + t.Fatal("tasks items must not have a type key (oneOf only)") + } + oneOf, ok := tasks["items"].(map[string]interface{})["oneOf"].([]interface{}) + if !ok || len(oneOf) != 2 { + t.Fatalf("tasks items oneOf = %v, want 2 branches", tasks["items"]) + } + if oneOf[0].(map[string]interface{})["type"] != "string" { + t.Fatalf("oneOf[0] = %v, want string", oneOf[0]) + } + obj := oneOf[1].(map[string]interface{}) + if obj["type"] != "object" { + t.Fatalf("oneOf[1] type = %v, want object", obj["type"]) + } + objReq, ok := obj["required"].([]string) + if !ok || len(objReq) != 1 || objReq[0] != "prompt" { + t.Fatalf("oneOf[1] required = %v, want [prompt]", obj["required"]) + } + req, _ := MultiAgentTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "tasks" { + t.Fatalf("required = %v, want [tasks]", MultiAgentTool{}.Parameters()["required"]) + } +} + +func TestSchemaPropertyExtensions(t *testing.T) { + p := SchemaProperty{Type: "string", MaxLength: 10, MinLength: 1} + m := p.toMap() + if m["maxLength"] != 10 || m["minLength"] != 1 { + t.Fatalf("length bounds = %v", m) + } + a := SchemaProperty{Type: "array", Items: &SchemaProperty{Type: "string"}, MaxItems: 5, MinItems: 1} + am := a.toMap() + if am["maxItems"] != 5 || am["minItems"] != 1 { + t.Fatalf("item bounds = %v", am) + } + if am["items"].(map[string]interface{})["type"] != "string" { + t.Fatalf("items = %v", am["items"]) + } + o := SchemaProperty{OneOf: []SchemaProperty{{Type: "string"}, {Type: "integer"}}} + om := o.toMap() + if _, hasType := om["type"]; hasType { + t.Fatalf("typeless oneOf must omit type key: %v", om) + } + branches, ok := om["oneOf"].([]interface{}) + if !ok || len(branches) != 2 || branches[0].(map[string]interface{})["type"] != "string" { + t.Fatalf("oneOf = %v", om["oneOf"]) + } +} diff --git a/internal/tool/tool.go b/internal/tool/tool.go index f2f423a6..a55ead1b 100644 --- a/internal/tool/tool.go +++ b/internal/tool/tool.go @@ -69,6 +69,19 @@ type SchemaProperty struct { // hand-written schemas exactly. Minimum interface{} `json:"minimum,omitempty"` Maximum interface{} `json:"maximum,omitempty"` + // MaxLength/MinLength constrain string fields. MaxItems/MinItems constrain + // array fields. Plain ints (omitted when zero): zero is never a meaningful + // bound, so authors write literals like MaxLength: 2000. + MaxLength int `json:"maxLength,omitempty"` + MinLength int `json:"minLength,omitempty"` + MaxItems int `json:"maxItems,omitempty"` + MinItems int `json:"minItems,omitempty"` + // OneOf/AnyOf list alternative subschemas (e.g. array items that accept a + // string or an object). Each branch is a full SchemaProperty. A property + // that only carries OneOf/AnyOf leaves Type empty, and toMap omits the + // "type" key in that case to match hand-written schemas exactly. + OneOf []SchemaProperty `json:"oneOf,omitempty"` + AnyOf []SchemaProperty `json:"anyOf,omitempty"` } // SchemaProvider is an optional interface tools can implement to expose a typed @@ -98,7 +111,10 @@ func (s ToolSchema) ToJSONSchema() map[string]interface{} { } func (p SchemaProperty) toMap() map[string]interface{} { - m := map[string]interface{}{"type": p.Type} + m := map[string]interface{}{} + if p.Type != "" { + m["type"] = p.Type + } if p.Description != "" { m["description"] = p.Description } @@ -114,6 +130,32 @@ func (p SchemaProperty) toMap() map[string]interface{} { if p.Maximum != nil { m["maximum"] = p.Maximum } + if p.MaxLength != 0 { + m["maxLength"] = p.MaxLength + } + if p.MinLength != 0 { + m["minLength"] = p.MinLength + } + if p.MaxItems != 0 { + m["maxItems"] = p.MaxItems + } + if p.MinItems != 0 { + m["minItems"] = p.MinItems + } + if len(p.OneOf) > 0 { + branches := make([]interface{}, 0, len(p.OneOf)) + for _, b := range p.OneOf { + branches = append(branches, b.toMap()) + } + m["oneOf"] = branches + } + if len(p.AnyOf) > 0 { + branches := make([]interface{}, 0, len(p.AnyOf)) + for _, b := range p.AnyOf { + branches = append(branches, b.toMap()) + } + m["anyOf"] = branches + } if p.Items != nil { m["items"] = p.Items.toMap() } diff --git a/internal/tool/web_search.go b/internal/tool/web_search.go index ae0ba21b..20dded64 100644 --- a/internal/tool/web_search.go +++ b/internal/tool/web_search.go @@ -30,55 +30,46 @@ type searchResult struct { type WebSearchTool struct{} -func (WebSearchTool) Name() string { return "WebSearch" } +func (WebSearchTool) Name() string { return "WebSearch" } + +// WebSearchInput is the typed input for WebSearchTool. +type WebSearchInput struct { + Query string `json:"query"` + Queries []string `json:"queries"` + NumResults int `json:"numResults"` + SearchType string `json:"searchType"` +} + func (WebSearchTool) RiskLevel() string { return "low" } func (WebSearchTool) Aliases() []string { return []string{"web_search"} } func (WebSearchTool) Description() string { return "Search the web and return structured results. Supports Brave Search, SearXNG, DeepSeek, Exa, Perplexity, and DuckDuckGo backends." } -func (WebSearchTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "query": map[string]interface{}{ - "type": "string", - "description": "Search query. Provide this OR queries (not both).", - "maxLength": maxWebSearchQueryLength, - }, - "queries": map[string]interface{}{ - "type": "array", - "items": map[string]interface{}{"type": "string", "maxLength": maxWebSearchQueryLength}, - "maxItems": maxWebSearchQueries, - "description": "Multiple search queries to run concurrently in a single call. Use this to research several things at once instead of issuing one WebSearch per query.", - }, - "numResults": map[string]interface{}{ - "type": "integer", - "description": "Number of results to return (1-20)", - "minimum": 1, - "maximum": 20, - "default": 5, - }, - "searchType": map[string]interface{}{ - "type": "string", - "description": "Type of search to perform", - "enum": []string{"web", "news"}, - "default": "web", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (WebSearchTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "query": {Type: "string", Description: "Search query. Provide this OR queries (not both).", MaxLength: maxWebSearchQueryLength}, + "queries": {Type: "array", Items: &SchemaProperty{Type: "string", MaxLength: maxWebSearchQueryLength}, MaxItems: maxWebSearchQueries, Description: "Multiple search queries to run concurrently in a single call. Use this to research several things at once instead of issuing one WebSearch per query."}, + "numResults": {Type: "integer", Description: "Number of results to return (1-20)", Minimum: 1, Maximum: 20, Default: 5}, + "searchType": {Type: "string", Description: "Type of search to perform", Enum: []interface{}{"web", "news"}, Default: "web"}, }, - // Either query or queries must be supplied; validated in Execute since - // JSON Schema "required" cannot express an exclusive-or cleanly. } } +func (WebSearchTool) Parameters() map[string]interface{} { + return webSearchSchema.ToJSONSchema() +} + +// webSearchSchema is the single source of truth for WebSearch's input schema. +var webSearchSchema = WebSearchTool{}.Schema() + func (t WebSearchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Query string `json:"query"` - Queries []string `json:"queries"` - NumResults int `json:"numResults"` - SearchType string `json:"searchType"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[WebSearchInput]("WebSearch", input) + if err != nil { return "", err }