1  import { z } from 'zod/v4'
  2  import type { TaskStateBase } from '../../Task.js'
  3  import { buildTool, type ToolDef } from '../../Tool.js'
  4  import { stopTask } from '../../tasks/stopTask.js'
  5  import { lazySchema } from '../../utils/lazySchema.js'
  6  import { jsonStringify } from '../../utils/slowOperations.js'
  7  import { DESCRIPTION, TASK_STOP_TOOL_NAME } from './prompt.js'
  8  import { renderToolResultMessage, renderToolUseMessage } from './UI.js'
  9  
 10  const inputSchema = lazySchema(() =>
 11    z.strictObject({
 12      task_id: z
 13        .string()
 14        .optional()
 15        .describe('The ID of the background task to stop'),
 16      // shell_id is accepted for backward compatibility with the deprecated KillShell tool
 17      shell_id: z.string().optional().describe('Deprecated: use task_id instead'),
 18    }),
 19  )
 20  type InputSchema = ReturnType<typeof inputSchema>
 21  
 22  const outputSchema = lazySchema(() =>
 23    z.object({
 24      message: z.string().describe('Status message about the operation'),
 25      task_id: z.string().describe('The ID of the task that was stopped'),
 26      task_type: z.string().describe('The type of the task that was stopped'),
 27      // Optional: tool outputs are persisted to transcripts and replayed on --resume
 28      // without re-validation, so sessions from before this field was added lack it.
 29      command: z
 30        .string()
 31        .optional()
 32        .describe('The command or description of the stopped task'),
 33    }),
 34  )
 35  type OutputSchema = ReturnType<typeof outputSchema>
 36  
 37  export type Output = z.infer<OutputSchema>
 38  
 39  export const TaskStopTool = buildTool({
 40    name: TASK_STOP_TOOL_NAME,
 41    searchHint: 'kill a running background task',
 42    // KillShell is the deprecated name - kept as alias for backward compatibility
 43    // with existing transcripts and SDK users
 44    aliases: ['KillShell'],
 45    maxResultSizeChars: 100_000,
 46    userFacingName: () => (process.env.USER_TYPE === 'ant' ? '' : 'Stop Task'),
 47    get inputSchema(): InputSchema {
 48      return inputSchema()
 49    },
 50    get outputSchema(): OutputSchema {
 51      return outputSchema()
 52    },
 53    shouldDefer: true,
 54    isConcurrencySafe() {
 55      return true
 56    },
 57    toAutoClassifierInput(input) {
 58      return input.task_id ?? input.shell_id ?? ''
 59    },
 60    async validateInput({ task_id, shell_id }, { getAppState }) {
 61      // Support both task_id and shell_id (deprecated KillShell compat)
 62      const id = task_id ?? shell_id
 63      if (!id) {
 64        return {
 65          result: false,
 66          message: 'Missing required parameter: task_id',
 67          errorCode: 1,
 68        }
 69      }
 70  
 71      const appState = getAppState()
 72      const task = appState.tasks?.[id] as TaskStateBase | undefined
 73  
 74      if (!task) {
 75        return {
 76          result: false,
 77          message: `No task found with ID: ${id}`,
 78          errorCode: 1,
 79        }
 80      }
 81  
 82      if (task.status !== 'running') {
 83        return {
 84          result: false,
 85          message: `Task ${id} is not running (status: ${task.status})`,
 86          errorCode: 3,
 87        }
 88      }
 89  
 90      return { result: true }
 91    },
 92    async description() {
 93      return `Stop a running background task by ID`
 94    },
 95    async prompt() {
 96      return DESCRIPTION
 97    },
 98    mapToolResultToToolResultBlockParam(output, toolUseID) {
 99      return {
100        tool_use_id: toolUseID,
101        type: 'tool_result',
102        content: jsonStringify(output),
103      }
104    },
105    renderToolUseMessage,
106    renderToolResultMessage,
107    async call(
108      { task_id, shell_id },
109      { getAppState, setAppState, abortController },
110    ) {
111      // Support both task_id and shell_id (deprecated KillShell compat)
112      const id = task_id ?? shell_id
113      if (!id) {
114        throw new Error('Missing required parameter: task_id')
115      }
116  
117      const result = await stopTask(id, {
118        getAppState,
119        setAppState,
120      })
121  
122      return {
123        data: {
124          message: `Successfully stopped task: ${result.taskId} (${result.command})`,
125          task_id: result.taskId,
126          task_type: result.taskType,
127          command: result.command,
128        },
129      }
130    },
131  } satisfies ToolDef<InputSchema, Output>)