Skip to content

Commit e286b44

Browse files
Patel230Patel230
authored
refactor(tool): migrate task tools to typed schemas (#332)
1 parent 60d3750 commit e286b44

5 files changed

Lines changed: 410 additions & 155 deletions

File tree

‎internal/tool/schema_batch_test.go‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1091,3 +1091,144 @@ func TestTicketComplianceSchemaProvider(t *testing.T) {
10911091
t.Fatalf("required = %v, want [ticket_content diff]", (&TicketComplianceTool{}).Parameters()["required"])
10921092
}
10931093
}
1094+
1095+
func TestWaitTasksSchemaProvider(t *testing.T) {
1096+
var _ SchemaProvider = WaitTasksTool{}
1097+
props := schemaProps(t, WaitTasksTool{}.Parameters())
1098+
if props["task_id"].(map[string]interface{})["type"] != "string" {
1099+
t.Fatal("task_id type wrong")
1100+
}
1101+
ids := props["task_ids"].(map[string]interface{})
1102+
if ids["type"] != "array" {
1103+
t.Fatalf("task_ids type = %v, want array", ids["type"])
1104+
}
1105+
if ids["items"].(map[string]interface{})["type"] != "string" {
1106+
t.Fatal("task_ids items type wrong")
1107+
}
1108+
if props["timeout_sec"].(map[string]interface{})["type"] != "integer" {
1109+
t.Fatal("timeout_sec type wrong")
1110+
}
1111+
}
1112+
1113+
func TestKillTaskSchemaProvider(t *testing.T) {
1114+
var _ SchemaProvider = KillTaskTool{}
1115+
props := schemaProps(t, KillTaskTool{}.Parameters())
1116+
if props["task_id"].(map[string]interface{})["type"] != "string" {
1117+
t.Fatal("task_id type wrong")
1118+
}
1119+
req, _ := KillTaskTool{}.Parameters()["required"].([]string)
1120+
if len(req) != 1 || req[0] != "task_id" {
1121+
t.Fatalf("required = %v, want [task_id]", KillTaskTool{}.Parameters()["required"])
1122+
}
1123+
}
1124+
1125+
func TestMonitorSchemaProvider(t *testing.T) {
1126+
var _ SchemaProvider = MonitorTool{}
1127+
props := schemaProps(t, MonitorTool{}.Parameters())
1128+
if props["command"].(map[string]interface{})["type"] != "string" {
1129+
t.Fatal("command type wrong")
1130+
}
1131+
if props["max_runtime_sec"].(map[string]interface{})["type"] != "integer" {
1132+
t.Fatal("max_runtime_sec type wrong")
1133+
}
1134+
if props["max_lines_per_sec"].(map[string]interface{})["type"] != "integer" {
1135+
t.Fatal("max_lines_per_sec type wrong")
1136+
}
1137+
req, _ := MonitorTool{}.Parameters()["required"].([]string)
1138+
if len(req) != 1 || req[0] != "command" {
1139+
t.Fatalf("required = %v, want [command]", MonitorTool{}.Parameters()["required"])
1140+
}
1141+
}
1142+
1143+
func TestTaskCreateSchemaProvider(t *testing.T) {
1144+
var _ SchemaProvider = TaskCreateTool{}
1145+
props := schemaProps(t, TaskCreateTool{}.Parameters())
1146+
if props["subject"].(map[string]interface{})["type"] != "string" {
1147+
t.Fatal("subject type wrong")
1148+
}
1149+
deps := props["dependencies"].(map[string]interface{})
1150+
if deps["type"] != "array" {
1151+
t.Fatalf("dependencies type = %v, want array", deps["type"])
1152+
}
1153+
depProps := deps["items"].(map[string]interface{})["properties"].(map[string]interface{})
1154+
depTypeEnum, ok := depProps["type"].(map[string]interface{})["enum"].([]interface{})
1155+
if !ok || len(depTypeEnum) != 3 || depTypeEnum[0] != "blocks" || depTypeEnum[2] != "parent-child" {
1156+
t.Fatalf("dependency type enum = %v, want 3 options", depProps["type"])
1157+
}
1158+
if props["metadata"].(map[string]interface{})["type"] != "object" {
1159+
t.Fatal("metadata type wrong")
1160+
}
1161+
req, _ := TaskCreateTool{}.Parameters()["required"].([]string)
1162+
if len(req) != 2 || req[0] != "subject" || req[1] != "description" {
1163+
t.Fatalf("required = %v, want [subject description]", TaskCreateTool{}.Parameters()["required"])
1164+
}
1165+
}
1166+
1167+
func TestTaskGetSchemaProvider(t *testing.T) {
1168+
var _ SchemaProvider = TaskGetTool{}
1169+
props := schemaProps(t, TaskGetTool{}.Parameters())
1170+
if props["taskId"].(map[string]interface{})["type"] != "string" {
1171+
t.Fatal("taskId type wrong")
1172+
}
1173+
req, _ := TaskGetTool{}.Parameters()["required"].([]string)
1174+
if len(req) != 1 || req[0] != "taskId" {
1175+
t.Fatalf("required = %v, want [taskId]", TaskGetTool{}.Parameters()["required"])
1176+
}
1177+
}
1178+
1179+
func TestTaskListSchemaProvider(t *testing.T) {
1180+
var _ SchemaProvider = TaskListTool{}
1181+
props := schemaProps(t, TaskListTool{}.Parameters())
1182+
enum, ok := props["action"].(map[string]interface{})["enum"].([]interface{})
1183+
if !ok || len(enum) != 4 || enum[0] != "list" || enum[3] != "compact" {
1184+
t.Fatalf("action enum = %v, want 4 options", props["action"])
1185+
}
1186+
}
1187+
1188+
func TestTaskUpdateSchemaProvider(t *testing.T) {
1189+
var _ SchemaProvider = TaskUpdateTool{}
1190+
props := schemaProps(t, TaskUpdateTool{}.Parameters())
1191+
enum, ok := props["status"].(map[string]interface{})["enum"].([]interface{})
1192+
if !ok || len(enum) != 7 || enum[0] != "pending" || enum[6] != "cancelled" {
1193+
t.Fatalf("status enum = %v, want 7 options", props["status"])
1194+
}
1195+
req, _ := TaskUpdateTool{}.Parameters()["required"].([]string)
1196+
if len(req) != 1 || req[0] != "taskId" {
1197+
t.Fatalf("required = %v, want [taskId]", TaskUpdateTool{}.Parameters()["required"])
1198+
}
1199+
}
1200+
1201+
func TestTaskRunSchemaProvider(t *testing.T) {
1202+
var _ SchemaProvider = TaskRunTool{}
1203+
props := schemaProps(t, TaskRunTool{}.Parameters())
1204+
if props["timeout_sec"].(map[string]interface{})["type"] != "integer" {
1205+
t.Fatal("timeout_sec type wrong")
1206+
}
1207+
if props["max_total_tasks"].(map[string]interface{})["type"] != "integer" {
1208+
t.Fatal("max_total_tasks type wrong")
1209+
}
1210+
}
1211+
1212+
func TestTaskOutputSchemaProvider(t *testing.T) {
1213+
var _ SchemaProvider = TaskOutputTool{}
1214+
props := schemaProps(t, TaskOutputTool{}.Parameters())
1215+
if props["task_id"].(map[string]interface{})["type"] != "string" {
1216+
t.Fatal("task_id type wrong")
1217+
}
1218+
req, _ := TaskOutputTool{}.Parameters()["required"].([]string)
1219+
if len(req) != 1 || req[0] != "task_id" {
1220+
t.Fatalf("required = %v, want [task_id]", TaskOutputTool{}.Parameters()["required"])
1221+
}
1222+
}
1223+
1224+
func TestTaskStopSchemaProvider(t *testing.T) {
1225+
var _ SchemaProvider = TaskStopTool{}
1226+
props := schemaProps(t, TaskStopTool{}.Parameters())
1227+
if props["task_id"].(map[string]interface{})["type"] != "string" {
1228+
t.Fatal("task_id type wrong")
1229+
}
1230+
req, _ := TaskStopTool{}.Parameters()["required"].([]string)
1231+
if len(req) != 1 || req[0] != "task_id" {
1232+
t.Fatalf("required = %v, want [task_id]", TaskStopTool{}.Parameters()["required"])
1233+
}
1234+
}

‎internal/tool/task_control.go‎

Lines changed: 75 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,13 @@ func killTask(ctx context.Context, id string) error {
7171
// WaitTasksTool waits for one or more background tasks.
7272
type WaitTasksTool struct{}
7373

74+
// WaitTasksInput is the typed input for WaitTasksTool.
75+
type WaitTasksInput struct {
76+
TaskID string `json:"task_id"`
77+
TaskIDs []string `json:"task_ids"`
78+
TimeoutSec int `json:"timeout_sec"`
79+
}
80+
7481
func (WaitTasksTool) Name() string { return "WaitTasks" }
7582
func (WaitTasksTool) Aliases() []string { return []string{"wait_tasks", "GetTaskOutput"} }
7683
func (WaitTasksTool) RiskLevel() string { return "low" }
@@ -79,34 +86,29 @@ func (WaitTasksTool) Description() string {
7986
"Works for shell (Bash run_in_background), agent spawns, and Monitor tasks."
8087
}
8188

82-
func (WaitTasksTool) Parameters() map[string]interface{} {
83-
return map[string]interface{}{
84-
"type": "object",
85-
"properties": map[string]interface{}{
86-
"task_id": map[string]interface{}{
87-
"type": "string",
88-
"description": "Single task ID to wait for",
89-
},
90-
"task_ids": map[string]interface{}{
91-
"type": "array",
92-
"items": map[string]interface{}{"type": "string"},
93-
"description": "Multiple task IDs to wait for",
94-
},
95-
"timeout_sec": map[string]interface{}{
96-
"type": "integer",
97-
"description": "Max seconds to wait (default 120, max 600)",
98-
},
89+
// Schema returns the typed input schema. Parameters() delegates to it so the
90+
// two cannot diverge.
91+
func (WaitTasksTool) Schema() ToolSchema {
92+
return ToolSchema{
93+
Type: "object",
94+
Properties: map[string]SchemaProperty{
95+
"task_id": {Type: "string", Description: "Single task ID to wait for"},
96+
"task_ids": {Type: "array", Items: &SchemaProperty{Type: "string"}, Description: "Multiple task IDs to wait for"},
97+
"timeout_sec": {Type: "integer", Description: "Max seconds to wait (default 120, max 600)"},
9998
},
10099
}
101100
}
102101

102+
func (WaitTasksTool) Parameters() map[string]interface{} {
103+
return waitTasksSchema.ToJSONSchema()
104+
}
105+
106+
// waitTasksSchema is the single source of truth for WaitTasks' input schema.
107+
var waitTasksSchema = WaitTasksTool{}.Schema()
108+
103109
func (WaitTasksTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
104-
var p struct {
105-
TaskID string `json:"task_id"`
106-
TaskIDs []string `json:"task_ids"`
107-
TimeoutSec int `json:"timeout_sec"`
108-
}
109-
if err := json.Unmarshal(input, &p); err != nil {
110+
p, err := DecodeInput[WaitTasksInput]("WaitTasks", input)
111+
if err != nil {
110112
return "", err
111113
}
112114
ids := append([]string{}, p.TaskIDs...)
@@ -165,28 +167,40 @@ func (WaitTasksTool) Execute(ctx context.Context, input json.RawMessage) (string
165167
// KillTaskTool stops a running background task.
166168
type KillTaskTool struct{}
167169

170+
// KillTaskInput is the typed input for KillTaskTool.
171+
type KillTaskInput struct {
172+
TaskID string `json:"task_id"`
173+
}
174+
168175
func (KillTaskTool) Name() string { return "KillTask" }
169176
func (KillTaskTool) Aliases() []string { return []string{"kill_task"} }
170177
func (KillTaskTool) RiskLevel() string { return "medium" }
171178
func (KillTaskTool) Description() string {
172179
return "Kill a running background shell, agent, or monitor task by task_id."
173180
}
174181

175-
func (KillTaskTool) Parameters() map[string]interface{} {
176-
return map[string]interface{}{
177-
"type": "object",
178-
"properties": map[string]interface{}{
179-
"task_id": map[string]interface{}{"type": "string", "description": "Background task ID"},
182+
// Schema returns the typed input schema. Parameters() delegates to it so the
183+
// two cannot diverge.
184+
func (KillTaskTool) Schema() ToolSchema {
185+
return ToolSchema{
186+
Type: "object",
187+
Properties: map[string]SchemaProperty{
188+
"task_id": {Type: "string", Description: "Background task ID"},
180189
},
181-
"required": []string{"task_id"},
190+
Required: []string{"task_id"},
182191
}
183192
}
184193

194+
func (KillTaskTool) Parameters() map[string]interface{} {
195+
return killTaskSchema.ToJSONSchema()
196+
}
197+
198+
// killTaskSchema is the single source of truth for KillTask's input schema.
199+
var killTaskSchema = KillTaskTool{}.Schema()
200+
185201
func (KillTaskTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
186-
var p struct {
187-
TaskID string `json:"task_id"`
188-
}
189-
if err := json.Unmarshal(input, &p); err != nil {
202+
p, err := DecodeInput[KillTaskInput]("KillTask", input)
203+
if err != nil {
190204
return "", err
191205
}
192206
if p.TaskID == "" {
@@ -201,6 +215,14 @@ func (KillTaskTool) Execute(ctx context.Context, input json.RawMessage) (string,
201215
// MonitorTool runs a command, streams lines into a background task, with rate limits.
202216
type MonitorTool struct{}
203217

218+
// MonitorInput is the typed input for MonitorTool.
219+
type MonitorInput struct {
220+
Command string `json:"command"`
221+
MaxRuntimeSec int `json:"max_runtime_sec"`
222+
MaxLinesPerSec int `json:"max_lines_per_sec"`
223+
Description string `json:"description"`
224+
}
225+
204226
func (MonitorTool) Name() string { return "Monitor" }
205227
func (MonitorTool) Aliases() []string { return []string{"monitor"} }
206228
func (MonitorTool) RiskLevel() string { return "medium" }
@@ -210,39 +232,31 @@ func (MonitorTool) Description() string {
210232
"Rate-limited (max lines/sec) and auto-killed after max_runtime_sec."
211233
}
212234

213-
func (MonitorTool) Parameters() map[string]interface{} {
214-
return map[string]interface{}{
215-
"type": "object",
216-
"properties": map[string]interface{}{
217-
"command": map[string]interface{}{
218-
"type": "string",
219-
"description": "Shell command to monitor (e.g. 'tail -f log' or 'npm run dev')",
220-
},
221-
"max_runtime_sec": map[string]interface{}{
222-
"type": "integer",
223-
"description": "Auto-kill after this many seconds (default 300, max 3600)",
224-
},
225-
"max_lines_per_sec": map[string]interface{}{
226-
"type": "integer",
227-
"description": "Drop excess lines above this rate (default 50)",
228-
},
229-
"description": map[string]interface{}{
230-
"type": "string",
231-
"description": "Short label for the monitor task",
232-
},
235+
// Schema returns the typed input schema. Parameters() delegates to it so the
236+
// two cannot diverge.
237+
func (MonitorTool) Schema() ToolSchema {
238+
return ToolSchema{
239+
Type: "object",
240+
Properties: map[string]SchemaProperty{
241+
"command": {Type: "string", Description: "Shell command to monitor (e.g. 'tail -f log' or 'npm run dev')"},
242+
"max_runtime_sec": {Type: "integer", Description: "Auto-kill after this many seconds (default 300, max 3600)"},
243+
"max_lines_per_sec": {Type: "integer", Description: "Drop excess lines above this rate (default 50)"},
244+
"description": {Type: "string", Description: "Short label for the monitor task"},
233245
},
234-
"required": []string{"command"},
246+
Required: []string{"command"},
235247
}
236248
}
237249

250+
func (MonitorTool) Parameters() map[string]interface{} {
251+
return monitorSchema.ToJSONSchema()
252+
}
253+
254+
// monitorSchema is the single source of truth for Monitor's input schema.
255+
var monitorSchema = MonitorTool{}.Schema()
256+
238257
func (MonitorTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
239-
var p struct {
240-
Command string `json:"command"`
241-
MaxRuntimeSec int `json:"max_runtime_sec"`
242-
MaxLinesPerSec int `json:"max_lines_per_sec"`
243-
Description string `json:"description"`
244-
}
245-
if err := json.Unmarshal(input, &p); err != nil {
258+
p, err := DecodeInput[MonitorInput]("Monitor", input)
259+
if err != nil {
246260
return "", err
247261
}
248262
if strings.TrimSpace(p.Command) == "" {

0 commit comments

Comments
 (0)