diff --git a/internal/tool/schema_batch_test.go b/internal/tool/schema_batch_test.go index 98215918..515584a3 100644 --- a/internal/tool/schema_batch_test.go +++ b/internal/tool/schema_batch_test.go @@ -1091,3 +1091,144 @@ func TestTicketComplianceSchemaProvider(t *testing.T) { t.Fatalf("required = %v, want [ticket_content diff]", (&TicketComplianceTool{}).Parameters()["required"]) } } + +func TestWaitTasksSchemaProvider(t *testing.T) { + var _ SchemaProvider = WaitTasksTool{} + props := schemaProps(t, WaitTasksTool{}.Parameters()) + if props["task_id"].(map[string]interface{})["type"] != "string" { + t.Fatal("task_id type wrong") + } + ids := props["task_ids"].(map[string]interface{}) + if ids["type"] != "array" { + t.Fatalf("task_ids type = %v, want array", ids["type"]) + } + if ids["items"].(map[string]interface{})["type"] != "string" { + t.Fatal("task_ids items type wrong") + } + if props["timeout_sec"].(map[string]interface{})["type"] != "integer" { + t.Fatal("timeout_sec type wrong") + } +} + +func TestKillTaskSchemaProvider(t *testing.T) { + var _ SchemaProvider = KillTaskTool{} + props := schemaProps(t, KillTaskTool{}.Parameters()) + if props["task_id"].(map[string]interface{})["type"] != "string" { + t.Fatal("task_id type wrong") + } + req, _ := KillTaskTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "task_id" { + t.Fatalf("required = %v, want [task_id]", KillTaskTool{}.Parameters()["required"]) + } +} + +func TestMonitorSchemaProvider(t *testing.T) { + var _ SchemaProvider = MonitorTool{} + props := schemaProps(t, MonitorTool{}.Parameters()) + if props["command"].(map[string]interface{})["type"] != "string" { + t.Fatal("command type wrong") + } + if props["max_runtime_sec"].(map[string]interface{})["type"] != "integer" { + t.Fatal("max_runtime_sec type wrong") + } + if props["max_lines_per_sec"].(map[string]interface{})["type"] != "integer" { + t.Fatal("max_lines_per_sec type wrong") + } + req, _ := MonitorTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "command" { + t.Fatalf("required = %v, want [command]", MonitorTool{}.Parameters()["required"]) + } +} + +func TestTaskCreateSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskCreateTool{} + props := schemaProps(t, TaskCreateTool{}.Parameters()) + if props["subject"].(map[string]interface{})["type"] != "string" { + t.Fatal("subject type wrong") + } + deps := props["dependencies"].(map[string]interface{}) + if deps["type"] != "array" { + t.Fatalf("dependencies type = %v, want array", deps["type"]) + } + depProps := deps["items"].(map[string]interface{})["properties"].(map[string]interface{}) + depTypeEnum, ok := depProps["type"].(map[string]interface{})["enum"].([]interface{}) + if !ok || len(depTypeEnum) != 3 || depTypeEnum[0] != "blocks" || depTypeEnum[2] != "parent-child" { + t.Fatalf("dependency type enum = %v, want 3 options", depProps["type"]) + } + if props["metadata"].(map[string]interface{})["type"] != "object" { + t.Fatal("metadata type wrong") + } + req, _ := TaskCreateTool{}.Parameters()["required"].([]string) + if len(req) != 2 || req[0] != "subject" || req[1] != "description" { + t.Fatalf("required = %v, want [subject description]", TaskCreateTool{}.Parameters()["required"]) + } +} + +func TestTaskGetSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskGetTool{} + props := schemaProps(t, TaskGetTool{}.Parameters()) + if props["taskId"].(map[string]interface{})["type"] != "string" { + t.Fatal("taskId type wrong") + } + req, _ := TaskGetTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "taskId" { + t.Fatalf("required = %v, want [taskId]", TaskGetTool{}.Parameters()["required"]) + } +} + +func TestTaskListSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskListTool{} + props := schemaProps(t, TaskListTool{}.Parameters()) + enum, ok := props["action"].(map[string]interface{})["enum"].([]interface{}) + if !ok || len(enum) != 4 || enum[0] != "list" || enum[3] != "compact" { + t.Fatalf("action enum = %v, want 4 options", props["action"]) + } +} + +func TestTaskUpdateSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskUpdateTool{} + props := schemaProps(t, TaskUpdateTool{}.Parameters()) + enum, ok := props["status"].(map[string]interface{})["enum"].([]interface{}) + if !ok || len(enum) != 7 || enum[0] != "pending" || enum[6] != "cancelled" { + t.Fatalf("status enum = %v, want 7 options", props["status"]) + } + req, _ := TaskUpdateTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "taskId" { + t.Fatalf("required = %v, want [taskId]", TaskUpdateTool{}.Parameters()["required"]) + } +} + +func TestTaskRunSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskRunTool{} + props := schemaProps(t, TaskRunTool{}.Parameters()) + if props["timeout_sec"].(map[string]interface{})["type"] != "integer" { + t.Fatal("timeout_sec type wrong") + } + if props["max_total_tasks"].(map[string]interface{})["type"] != "integer" { + t.Fatal("max_total_tasks type wrong") + } +} + +func TestTaskOutputSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskOutputTool{} + props := schemaProps(t, TaskOutputTool{}.Parameters()) + if props["task_id"].(map[string]interface{})["type"] != "string" { + t.Fatal("task_id type wrong") + } + req, _ := TaskOutputTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "task_id" { + t.Fatalf("required = %v, want [task_id]", TaskOutputTool{}.Parameters()["required"]) + } +} + +func TestTaskStopSchemaProvider(t *testing.T) { + var _ SchemaProvider = TaskStopTool{} + props := schemaProps(t, TaskStopTool{}.Parameters()) + if props["task_id"].(map[string]interface{})["type"] != "string" { + t.Fatal("task_id type wrong") + } + req, _ := TaskStopTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "task_id" { + t.Fatalf("required = %v, want [task_id]", TaskStopTool{}.Parameters()["required"]) + } +} diff --git a/internal/tool/task_control.go b/internal/tool/task_control.go index ae52da69..e37c85d5 100644 --- a/internal/tool/task_control.go +++ b/internal/tool/task_control.go @@ -71,6 +71,13 @@ func killTask(ctx context.Context, id string) error { // WaitTasksTool waits for one or more background tasks. type WaitTasksTool struct{} +// WaitTasksInput is the typed input for WaitTasksTool. +type WaitTasksInput struct { + TaskID string `json:"task_id"` + TaskIDs []string `json:"task_ids"` + TimeoutSec int `json:"timeout_sec"` +} + func (WaitTasksTool) Name() string { return "WaitTasks" } func (WaitTasksTool) Aliases() []string { return []string{"wait_tasks", "GetTaskOutput"} } func (WaitTasksTool) RiskLevel() string { return "low" } @@ -79,34 +86,29 @@ func (WaitTasksTool) Description() string { "Works for shell (Bash run_in_background), agent spawns, and Monitor tasks." } -func (WaitTasksTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "task_id": map[string]interface{}{ - "type": "string", - "description": "Single task ID to wait for", - }, - "task_ids": map[string]interface{}{ - "type": "array", - "items": map[string]interface{}{"type": "string"}, - "description": "Multiple task IDs to wait for", - }, - "timeout_sec": map[string]interface{}{ - "type": "integer", - "description": "Max seconds to wait (default 120, max 600)", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (WaitTasksTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "task_id": {Type: "string", Description: "Single task ID to wait for"}, + "task_ids": {Type: "array", Items: &SchemaProperty{Type: "string"}, Description: "Multiple task IDs to wait for"}, + "timeout_sec": {Type: "integer", Description: "Max seconds to wait (default 120, max 600)"}, }, } } +func (WaitTasksTool) Parameters() map[string]interface{} { + return waitTasksSchema.ToJSONSchema() +} + +// waitTasksSchema is the single source of truth for WaitTasks' input schema. +var waitTasksSchema = WaitTasksTool{}.Schema() + func (WaitTasksTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - TaskID string `json:"task_id"` - TaskIDs []string `json:"task_ids"` - TimeoutSec int `json:"timeout_sec"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[WaitTasksInput]("WaitTasks", input) + if err != nil { return "", err } ids := append([]string{}, p.TaskIDs...) @@ -165,6 +167,11 @@ func (WaitTasksTool) Execute(ctx context.Context, input json.RawMessage) (string // KillTaskTool stops a running background task. type KillTaskTool struct{} +// KillTaskInput is the typed input for KillTaskTool. +type KillTaskInput struct { + TaskID string `json:"task_id"` +} + func (KillTaskTool) Name() string { return "KillTask" } func (KillTaskTool) Aliases() []string { return []string{"kill_task"} } func (KillTaskTool) RiskLevel() string { return "medium" } @@ -172,21 +179,28 @@ func (KillTaskTool) Description() string { return "Kill a running background shell, agent, or monitor task by task_id." } -func (KillTaskTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "task_id": map[string]interface{}{"type": "string", "description": "Background task ID"}, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (KillTaskTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "task_id": {Type: "string", Description: "Background task ID"}, }, - "required": []string{"task_id"}, + Required: []string{"task_id"}, } } +func (KillTaskTool) Parameters() map[string]interface{} { + return killTaskSchema.ToJSONSchema() +} + +// killTaskSchema is the single source of truth for KillTask's input schema. +var killTaskSchema = KillTaskTool{}.Schema() + func (KillTaskTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - TaskID string `json:"task_id"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[KillTaskInput]("KillTask", input) + if err != nil { return "", err } if p.TaskID == "" { @@ -201,6 +215,14 @@ func (KillTaskTool) Execute(ctx context.Context, input json.RawMessage) (string, // MonitorTool runs a command, streams lines into a background task, with rate limits. type MonitorTool struct{} +// MonitorInput is the typed input for MonitorTool. +type MonitorInput struct { + Command string `json:"command"` + MaxRuntimeSec int `json:"max_runtime_sec"` + MaxLinesPerSec int `json:"max_lines_per_sec"` + Description string `json:"description"` +} + func (MonitorTool) Name() string { return "Monitor" } func (MonitorTool) Aliases() []string { return []string{"monitor"} } func (MonitorTool) RiskLevel() string { return "medium" } @@ -210,39 +232,31 @@ func (MonitorTool) Description() string { "Rate-limited (max lines/sec) and auto-killed after max_runtime_sec." } -func (MonitorTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "command": map[string]interface{}{ - "type": "string", - "description": "Shell command to monitor (e.g. 'tail -f log' or 'npm run dev')", - }, - "max_runtime_sec": map[string]interface{}{ - "type": "integer", - "description": "Auto-kill after this many seconds (default 300, max 3600)", - }, - "max_lines_per_sec": map[string]interface{}{ - "type": "integer", - "description": "Drop excess lines above this rate (default 50)", - }, - "description": map[string]interface{}{ - "type": "string", - "description": "Short label for the monitor task", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (MonitorTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "command": {Type: "string", Description: "Shell command to monitor (e.g. 'tail -f log' or 'npm run dev')"}, + "max_runtime_sec": {Type: "integer", Description: "Auto-kill after this many seconds (default 300, max 3600)"}, + "max_lines_per_sec": {Type: "integer", Description: "Drop excess lines above this rate (default 50)"}, + "description": {Type: "string", Description: "Short label for the monitor task"}, }, - "required": []string{"command"}, + Required: []string{"command"}, } } +func (MonitorTool) Parameters() map[string]interface{} { + return monitorSchema.ToJSONSchema() +} + +// monitorSchema is the single source of truth for Monitor's input schema. +var monitorSchema = MonitorTool{}.Schema() + func (MonitorTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Command string `json:"command"` - MaxRuntimeSec int `json:"max_runtime_sec"` - MaxLinesPerSec int `json:"max_lines_per_sec"` - Description string `json:"description"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[MonitorInput]("Monitor", input) + if err != nil { return "", err } if strings.TrimSpace(p.Command) == "" { diff --git a/internal/tool/task_create.go b/internal/tool/task_create.go index 3c4bbd9b..199bfece 100644 --- a/internal/tool/task_create.go +++ b/internal/tool/task_create.go @@ -212,34 +212,53 @@ func (s *TaskStore) CompactCompleted() string { // TaskCreateTool creates a new task in the task list. type TaskCreateTool struct{} +// TaskCreateInput is the typed input for TaskCreateTool. +type TaskCreateInput struct { + Subject string `json:"subject"` + Description string `json:"description"` + ActiveForm string `json:"activeForm"` + ParentID string `json:"parentId"` + Dependencies []TaskDependency `json:"dependencies"` + Metadata map[string]any `json:"metadata"` +} + func (TaskCreateTool) Name() string { return "TaskCreate" } func (TaskCreateTool) Aliases() []string { return []string{"task_create"} } func (TaskCreateTool) Description() string { return "Create a new task in the task list" } -func (TaskCreateTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "subject": map[string]interface{}{"type": "string", "description": "A brief title for the task"}, - "description": map[string]interface{}{"type": "string", "description": "What needs to be done"}, - "activeForm": map[string]interface{}{"type": "string", "description": "Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")"}, - "parentId": map[string]interface{}{"type": "string", "description": "Parent task ID for hierarchical tasks"}, - "dependencies": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "object", "properties": map[string]interface{}{"targetId": map[string]interface{}{"type": "string"}, "type": map[string]interface{}{"type": "string", "enum": []string{"blocks", "related", "parent-child"}}}}, "description": "Typed dependencies"}, - "metadata": map[string]interface{}{"type": "object", "description": "Arbitrary metadata to attach to the task"}, + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskCreateTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "subject": {Type: "string", Description: "A brief title for the task"}, + "description": {Type: "string", Description: "What needs to be done"}, + "activeForm": {Type: "string", Description: "Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")"}, + "parentId": {Type: "string", Description: "Parent task ID for hierarchical tasks"}, + "dependencies": {Type: "array", Items: &SchemaProperty{ + Type: "object", + Properties: map[string]SchemaProperty{ + "targetId": {Type: "string"}, + "type": {Type: "string", Enum: []interface{}{"blocks", "related", "parent-child"}}, + }, + }, Description: "Typed dependencies"}, + "metadata": {Type: "object", Description: "Arbitrary metadata to attach to the task"}, }, - "required": []string{"subject", "description"}, + Required: []string{"subject", "description"}, } } +func (TaskCreateTool) Parameters() map[string]interface{} { + return taskCreateSchema.ToJSONSchema() +} + +// taskCreateSchema is the single source of truth for TaskCreate's input schema. +var taskCreateSchema = TaskCreateTool{}.Schema() + func (TaskCreateTool) Execute(_ context.Context, input json.RawMessage) (string, error) { - var p struct { - Subject string `json:"subject"` - Description string `json:"description"` - ActiveForm string `json:"activeForm"` - ParentID string `json:"parentId"` - Dependencies []TaskDependency `json:"dependencies"` - Metadata map[string]any `json:"metadata"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[TaskCreateInput]("TaskCreate", input) + if err != nil { return "", err } if p.Subject == "" { @@ -263,24 +282,37 @@ func (TaskCreateTool) Execute(_ context.Context, input json.RawMessage) (string, // TaskGetTool retrieves a task by ID. type TaskGetTool struct{} +// TaskGetInput is the typed input for TaskGetTool. +type TaskGetInput struct { + TaskID string `json:"taskId"` +} + func (TaskGetTool) Name() string { return "TaskGet" } func (TaskGetTool) Aliases() []string { return []string{"task_get"} } func (TaskGetTool) Description() string { return "Get a task by ID from the task list" } -func (TaskGetTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "taskId": map[string]interface{}{"type": "string", "description": "The ID of the task to retrieve"}, + +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskGetTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "taskId": {Type: "string", Description: "The ID of the task to retrieve"}, }, - "required": []string{"taskId"}, + Required: []string{"taskId"}, } } +func (TaskGetTool) Parameters() map[string]interface{} { + return taskGetSchema.ToJSONSchema() +} + +// taskGetSchema is the single source of truth for TaskGet's input schema. +var taskGetSchema = TaskGetTool{}.Schema() + func (TaskGetTool) Execute(_ context.Context, input json.RawMessage) (string, error) { - var p struct { - TaskID string `json:"taskId"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[TaskGetInput]("TaskGet", input) + if err != nil { return "", err } task, ok := globalTaskStore.Get(p.TaskID) @@ -313,25 +345,41 @@ func (TaskGetTool) Execute(_ context.Context, input json.RawMessage) (string, er // TaskListTool lists all tasks. type TaskListTool struct{} +// TaskListInput is the typed input for TaskListTool. +type TaskListInput struct { + Action string `json:"action"` +} + func (TaskListTool) Name() string { return "TaskList" } func (TaskListTool) Aliases() []string { return []string{"task_list"} } func (TaskListTool) Description() string { return "List all tasks, ready tasks, or compact completed" } -func (TaskListTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "action": map[string]interface{}{"type": "string", "enum": []string{"list", "ready", "failed", "compact"}, "description": "Action: list (default), ready (pending with no blockers), failed (replan candidates), compact (remove completed)"}, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskListTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "action": {Type: "string", Enum: []interface{}{"list", "ready", "failed", "compact"}, Description: "Action: list (default), ready (pending with no blockers), failed (replan candidates), compact (remove completed)"}, }, } } +func (TaskListTool) Parameters() map[string]interface{} { + return taskListSchema.ToJSONSchema() +} + +// taskListSchema is the single source of truth for TaskList's input schema. +var taskListSchema = TaskListTool{}.Schema() + func (TaskListTool) Execute(_ context.Context, input json.RawMessage) (string, error) { - var p struct { - Action string `json:"action"` - } - if input != nil { - _ = json.Unmarshal(input, &p) + var p TaskListInput + if len(input) > 0 && string(input) != "null" { + decoded, err := DecodeInput[TaskListInput]("TaskList", input) + if err != nil { + return "", err + } + p = decoded } switch p.Action { @@ -393,31 +441,49 @@ func (TaskListTool) Execute(_ context.Context, input json.RawMessage) (string, e // TaskUpdateTool updates task fields. type TaskUpdateTool struct{} +// TaskUpdateInput is the typed input for TaskUpdateTool. +type TaskUpdateInput struct { + TaskID string `json:"taskId"` + Status string `json:"status"` + Owner string `json:"owner"` + Dependencies []TaskDependency `json:"dependencies"` +} + func (TaskUpdateTool) Name() string { return "TaskUpdate" } func (TaskUpdateTool) Aliases() []string { return []string{"task_update"} } func (TaskUpdateTool) Description() string { return "Update a task's status, owner, or dependencies" } -func (TaskUpdateTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "taskId": map[string]interface{}{"type": "string", "description": "The ID of the task to update"}, - "status": map[string]interface{}{"type": "string", "enum": []string{"pending", "in_progress", "reviewing", "completed", "failed", "skipped", "cancelled"}, "description": "New task status"}, - "owner": map[string]interface{}{"type": "string", "description": "Agent name to assign"}, - "dependencies": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "object", "properties": map[string]interface{}{"targetId": map[string]interface{}{"type": "string"}, "type": map[string]interface{}{"type": "string", "enum": []string{"blocks", "related", "parent-child"}}}}, "description": "Replace dependencies"}, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskUpdateTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "taskId": {Type: "string", Description: "The ID of the task to update"}, + "status": {Type: "string", Enum: []interface{}{"pending", "in_progress", "reviewing", "completed", "failed", "skipped", "cancelled"}, Description: "New task status"}, + "owner": {Type: "string", Description: "Agent name to assign"}, + "dependencies": {Type: "array", Items: &SchemaProperty{ + Type: "object", + Properties: map[string]SchemaProperty{ + "targetId": {Type: "string"}, + "type": {Type: "string", Enum: []interface{}{"blocks", "related", "parent-child"}}, + }, + }, Description: "Replace dependencies"}, }, - "required": []string{"taskId"}, + Required: []string{"taskId"}, } } +func (TaskUpdateTool) Parameters() map[string]interface{} { + return taskUpdateSchema.ToJSONSchema() +} + +// taskUpdateSchema is the single source of truth for TaskUpdate's input schema. +var taskUpdateSchema = TaskUpdateTool{}.Schema() + func (TaskUpdateTool) Execute(_ context.Context, input json.RawMessage) (string, error) { - var p struct { - TaskID string `json:"taskId"` - Status string `json:"status"` - Owner string `json:"owner"` - Dependencies []TaskDependency `json:"dependencies"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[TaskUpdateInput]("TaskUpdate", input) + if err != nil { return "", err } if p.TaskID == "" { diff --git a/internal/tool/task_run_tool.go b/internal/tool/task_run_tool.go index ed5ccf58..db8b0c08 100644 --- a/internal/tool/task_run_tool.go +++ b/internal/tool/task_run_tool.go @@ -15,6 +15,12 @@ import ( // for TaskRunner. type TaskRunTool struct{} +// TaskRunInput is the typed input for TaskRunTool. +type TaskRunInput struct { + TimeoutSec int `json:"timeout_sec"` + MaxTotalTasks int `json:"max_total_tasks"` +} + func (TaskRunTool) Name() string { return "TaskRun" } func (TaskRunTool) Aliases() []string { return []string{"task_run"} } func (TaskRunTool) Description() string { @@ -23,34 +29,38 @@ func (TaskRunTool) Description() string { "when the budget is exhausted. Returns a run summary." } -func (TaskRunTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "timeout_sec": map[string]interface{}{ - "type": "integer", - "description": "Per-task execution timeout in seconds (default 300)", - }, - "max_total_tasks": map[string]interface{}{ - "type": "integer", - "description": "Watchdog cap on distinct tasks that may reach a terminal state (default 50)", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskRunTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "timeout_sec": {Type: "integer", Description: "Per-task execution timeout in seconds (default 300)"}, + "max_total_tasks": {Type: "integer", Description: "Watchdog cap on distinct tasks that may reach a terminal state (default 50)"}, }, } } +func (TaskRunTool) Parameters() map[string]interface{} { + return taskRunSchema.ToJSONSchema() +} + +// taskRunSchema is the single source of truth for TaskRun's input schema. +var taskRunSchema = TaskRunTool{}.Schema() + func (TaskRunTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { tc := GetToolContext(ctx) if tc == nil || tc.TaskExecutor == nil { return "", fmt.Errorf("TaskRun requires a task executor; none is configured for this session") } - var p struct { - TimeoutSec int `json:"timeout_sec"` - MaxTotalTasks int `json:"max_total_tasks"` - } - if input != nil { - _ = json.Unmarshal(input, &p) + var p TaskRunInput + if len(input) > 0 && string(input) != "null" { + decoded, err := DecodeInput[TaskRunInput]("TaskRun", input) + if err != nil { + return "", err + } + p = decoded } timeout := time.Duration(p.TimeoutSec) * time.Second if timeout <= 0 { diff --git a/internal/tool/task_tools.go b/internal/tool/task_tools.go index 0b315e69..e431f2ac 100644 --- a/internal/tool/task_tools.go +++ b/internal/tool/task_tools.go @@ -231,27 +231,39 @@ func (t *backgroundTask) stop() error { type TaskOutputTool struct{} +// TaskOutputInput is the typed input for TaskOutputTool. +type TaskOutputInput struct { + TaskID string `json:"task_id"` +} + func (TaskOutputTool) Name() string { return "TaskOutput" } func (TaskOutputTool) Aliases() []string { return []string{"task_output"} } func (TaskOutputTool) Description() string { return "Read output from a background task (shell, agent, or monitor) by task_id." } -func (TaskOutputTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "task_id": map[string]interface{}{"type": "string", "description": "Background task ID"}, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskOutputTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "task_id": {Type: "string", Description: "Background task ID"}, }, - "required": []string{"task_id"}, + Required: []string{"task_id"}, } } +func (TaskOutputTool) Parameters() map[string]interface{} { + return taskOutputSchema.ToJSONSchema() +} + +// taskOutputSchema is the single source of truth for TaskOutput's input schema. +var taskOutputSchema = TaskOutputTool{}.Schema() + func (TaskOutputTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - TaskID string `json:"task_id"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[TaskOutputInput]("TaskOutput", input) + if err != nil { return "", err } // Accept legacy taskId field @@ -271,27 +283,39 @@ func (TaskOutputTool) Execute(ctx context.Context, input json.RawMessage) (strin type TaskStopTool struct{} +// TaskStopInput is the typed input for TaskStopTool. +type TaskStopInput struct { + TaskID string `json:"task_id"` +} + func (TaskStopTool) Name() string { return "TaskStop" } func (TaskStopTool) Aliases() []string { return []string{"task_stop"} } func (TaskStopTool) Description() string { return "Stop a background shell, agent, or monitor task." } -func (TaskStopTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "task_id": map[string]interface{}{"type": "string", "description": "Background task ID"}, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (TaskStopTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "task_id": {Type: "string", Description: "Background task ID"}, }, - "required": []string{"task_id"}, + Required: []string{"task_id"}, } } +func (TaskStopTool) Parameters() map[string]interface{} { + return taskStopSchema.ToJSONSchema() +} + +// taskStopSchema is the single source of truth for TaskStop's input schema. +var taskStopSchema = TaskStopTool{}.Schema() + func (TaskStopTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - TaskID string `json:"task_id"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[TaskStopInput]("TaskStop", input) + if err != nil { return "", err } if p.TaskID == "" {