Skip to content

refactor(tool): migrate batch12 task tools to typed schemas #332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Patel230 merged 1 commit into from
Sep 16, 2026
Merged

refactor(tool): migrate batch12 task tools to typed schemas #332

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 141 additions & 0 deletions internal/tool/schema_batch_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
}
}
136 changes: 75 additions & 61 deletions internal/tool/task_control.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
Expand All @@ -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...)
Expand Down Expand Up @@ -165,28 +167,40 @@ 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" }
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 == "" {
Expand All @@ -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" }
Expand All @@ -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) == "" {
Expand Down
Loading
Loading