diff --git a/packages/builtin-tools/src/index.ts b/packages/builtin-tools/src/index.ts index c31d600b33..d707826220 100644 --- a/packages/builtin-tools/src/index.ts +++ b/packages/builtin-tools/src/index.ts @@ -12,6 +12,7 @@ export { AskUserQuestionTool } from './tools/AskUserQuestionTool/AskUserQuestion export { BashTool } from './tools/BashTool/BashTool.js' export { BriefTool } from './tools/BriefTool/BriefTool.js' export { ConfigTool } from './tools/ConfigTool/ConfigTool.js' +export { GoalTool } from './tools/GoalTool/GoalTool.js' export { EnterPlanModeTool } from './tools/EnterPlanModeTool/EnterPlanModeTool.js' export { EnterWorktreeTool } from './tools/EnterWorktreeTool/EnterWorktreeTool.js' export { ExitPlanModeV2Tool } from './tools/ExitPlanModeTool/ExitPlanModeV2Tool.js' diff --git a/packages/builtin-tools/src/tools/GoalTool/GoalTool.ts b/packages/builtin-tools/src/tools/GoalTool/GoalTool.ts new file mode 100644 index 0000000000..2c754fd3e6 --- /dev/null +++ b/packages/builtin-tools/src/tools/GoalTool/GoalTool.ts @@ -0,0 +1,212 @@ +import { z } from 'zod/v4' +import { buildTool, type ToolDef } from 'src/Tool.js' +import { lazySchema } from 'src/utils/lazySchema.js' +import { + completeGoal, + formatGoalStatus, + getActiveElapsedMs, + getGoal, + setGoal, +} from 'src/services/goal/goalState.js' +import { DESCRIPTION, generatePrompt } from './prompt.js' +import type { ToolResultBlockParam } from '@anthropic-ai/sdk/resources/index.mjs' + +const inputSchema = lazySchema(() => + z.strictObject({ + action: z + .enum(['get', 'set', 'complete']) + .describe('The action to perform on the goal.'), + objective: z + .string() + .optional() + .describe('The goal objective. Required for "set" action.'), + message: z + .string() + .optional() + .describe('Completion message for "complete" action.'), + }), +) +type InputSchema = ReturnType + +const outputSchema = lazySchema(() => + z.object({ + success: z.boolean(), + action: z.string(), + goal: z + .object({ + objective: z.string(), + status: z.string(), + tokensUsed: z.number(), + tokenBudget: z.number().nullable(), + elapsedSeconds: z.number(), + }) + .optional(), + message: z.string().optional(), + error: z.string().optional(), + }), +) +type OutputSchema = ReturnType + +export type Input = z.infer +export type Output = z.infer + +export const GoalTool = buildTool({ + name: 'goal', + searchHint: 'manage long-running task goals', + maxResultSizeChars: 10_000, + async description() { + return DESCRIPTION + }, + async prompt() { + return generatePrompt() + }, + get inputSchema(): InputSchema { + return inputSchema() + }, + get outputSchema(): OutputSchema { + return outputSchema() + }, + userFacingName() { + return 'Goal' + }, + shouldDefer: true, + isConcurrencySafe() { + return true + }, + isReadOnly(input: Input) { + return input.action === 'get' + }, + toAutoClassifierInput(input) { + if (input.action === 'get') return 'get goal status' + if (input.action === 'set') return `set goal: ${input.objective}` + return `complete goal: ${input.message ?? ''}` + }, + async checkPermissions(input: Input) { + if (input.action === 'get') { + return { behavior: 'allow' as const, updatedInput: input } + } + return { + behavior: 'ask' as const, + message: + input.action === 'set' + ? `Set goal: ${input.objective}` + : `Complete goal${input.message ? `: ${input.message}` : ''}`, + } + }, + async call({ action, objective, message }: Input): Promise<{ data: Output }> { + if (action === 'get') { + const goal = getGoal() + if (!goal) { + return { data: { success: true, action, message: 'No active goal.' } } + } + const elapsedSeconds = Math.floor(getActiveElapsedMs(goal) / 1000) + return { + data: { + success: true, + action, + goal: { + objective: goal.objective, + status: goal.status, + tokensUsed: goal.tokensUsed, + tokenBudget: goal.tokenBudget, + elapsedSeconds, + }, + }, + } + } + + if (action === 'set') { + if (!objective) { + return { + data: { + success: false, + action, + error: 'objective is required for set action.', + }, + } + } + setGoal(objective) + return { + data: { + success: true, + action, + message: `Goal set: ${objective}`, + goal: { + objective, + status: 'active', + tokensUsed: 0, + tokenBudget: null, + elapsedSeconds: 0, + }, + }, + } + } + + if (action === 'complete') { + if (!completeGoal()) { + return { + data: { + success: false, + action, + error: 'No active goal to complete.', + }, + } + } + return { + data: { + success: true, + action, + message: message + ? `Goal completed: ${message}` + : 'Goal marked as complete.', + }, + } + } + + return { + data: { success: false, action, error: `Unknown action: ${action}` }, + } + }, + renderToolUseMessage(input: Partial) { + if (input.action === 'get') return 'Getting goal status' + if (input.action === 'set') return `Setting goal: ${input.objective ?? ''}` + if (input.action === 'complete') return 'Completing goal' + return 'Managing goal' + }, + renderToolResultMessage(content: Output) { + if (!content.success) return `Error: ${content.error}` + if (content.action === 'get' && content.goal) { + const g = content.goal + return `Goal: ${g.objective} [${g.status}]` + } + return content.message ?? 'Done.' + }, + mapToolResultToToolResultBlockParam( + content: Output, + toolUseID: string, + ): ToolResultBlockParam { + if (!content.success) { + return { + tool_use_id: toolUseID, + type: 'tool_result' as const, + content: `Error: ${content.error}`, + is_error: true, + } + } + + if (content.action === 'get' && content.goal) { + const g = content.goal + return { + tool_use_id: toolUseID, + type: 'tool_result' as const, + content: `Goal: ${g.objective}\nStatus: ${g.status}\nTokens: ${g.tokensUsed}${g.tokenBudget !== null ? ` / ${g.tokenBudget}` : ''}\nElapsed: ${g.elapsedSeconds}s`, + } + } + + return { + tool_use_id: toolUseID, + type: 'tool_result' as const, + content: content.message ?? 'Done.', + } + }, +} satisfies ToolDef) diff --git a/packages/builtin-tools/src/tools/GoalTool/prompt.ts b/packages/builtin-tools/src/tools/GoalTool/prompt.ts new file mode 100644 index 0000000000..a83fa7b669 --- /dev/null +++ b/packages/builtin-tools/src/tools/GoalTool/prompt.ts @@ -0,0 +1,18 @@ +export const DESCRIPTION = 'Manage the active goal for long-running tasks.' + +export function generatePrompt(): string { + return `Manage the active goal for long-running tasks. + +Use this tool to get, set, or complete a goal. A goal is an objective that the system tracks across the session, injecting continuation prompts to keep working toward it. + +## Actions +- **get** — Get the current goal status +- **set** — Set or update the goal objective +- **complete** — Mark the goal as complete when the objective is achieved + +## Examples +- Get current goal: { "action": "get" } +- Set a goal: { "action": "set", "objective": "Improve test coverage to 80%" } +- Complete a goal: { "action": "complete", "message": "All tests now pass with 82% coverage." } +` +} diff --git a/src/commands.ts b/src/commands.ts index 012a6a9bb0..af7da9ef8b 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -167,6 +167,7 @@ import thinkbackPlay from './commands/thinkback-play/index.js' import permissions from './commands/permissions/index.js' import plan from './commands/plan/index.js' import fast from './commands/fast/index.js' +import goal from './commands/goal/index.js' import passes from './commands/passes/index.js' import privacySettings from './commands/privacy-settings/index.js' import hooks from './commands/hooks/index.js' @@ -316,6 +317,7 @@ const COMMANDS = memoize((): Command[] => [ exit, fast, files, + goal, heapDump, help, ide, diff --git a/src/commands/goal/goal.ts b/src/commands/goal/goal.ts new file mode 100644 index 0000000000..cb639e8af6 --- /dev/null +++ b/src/commands/goal/goal.ts @@ -0,0 +1,66 @@ +import type { LocalCommandCall } from '../../types/command.js' +import { + clearGoal, + completeGoal, + formatGoalStatus, + getGoal, + pauseGoal, + resumeGoal, + setGoal, +} from '../../services/goal/goalState.js' + +export const call: LocalCommandCall = async args => { + const trimmed = args.trim() + + // No arguments — show current goal status + if (!trimmed) { + return { type: 'text', value: formatGoalStatus() } + } + + const lower = trimmed.toLowerCase() + + // Control subcommands + if (lower === 'clear') { + const goal = getGoal() + if (!goal) { + return { type: 'text', value: 'No active goal to clear.' } + } + clearGoal() + return { type: 'text', value: 'Goal cleared.' } + } + + if (lower === 'pause') { + if (pauseGoal()) { + return { type: 'text', value: 'Goal paused.' } + } + return { type: 'text', value: 'No active goal to pause.' } + } + + if (lower === 'resume') { + if (resumeGoal()) { + return { type: 'text', value: 'Goal resumed.' } + } + return { type: 'text', value: 'No paused goal to resume.' } + } + + if (lower === 'complete') { + if (completeGoal()) { + return { type: 'text', value: 'Goal marked as complete.' } + } + return { type: 'text', value: 'No active goal to complete.' } + } + + // Set a new goal + const existing = getGoal() + if (existing && existing.status === 'active') { + // Replace existing active goal + setGoal(trimmed) + return { + type: 'text', + value: `Goal replaced.\n\n${formatGoalStatus()}`, + } + } + + setGoal(trimmed) + return { type: 'text', value: `Goal set.\n\n${formatGoalStatus()}` } +} diff --git a/src/commands/goal/index.ts b/src/commands/goal/index.ts new file mode 100644 index 0000000000..7979eb0216 --- /dev/null +++ b/src/commands/goal/index.ts @@ -0,0 +1,12 @@ +import type { Command } from '../../commands.js' + +const goal = { + type: 'local', + name: 'goal', + description: 'Set or view the goal for a long-running task', + supportsNonInteractive: true, + argumentHint: ' | clear | pause | resume', + load: () => import('./goal.js'), +} satisfies Command + +export default goal diff --git a/src/constants/prompts.ts b/src/constants/prompts.ts index cca0a4264f..fad2b858ac 100644 --- a/src/constants/prompts.ts +++ b/src/constants/prompts.ts @@ -57,6 +57,7 @@ import { resolveSystemPromptSections, } from './systemPromptSections.js' import { SLEEP_TOOL_NAME } from '@claude-code-best/builtin-tools/tools/SleepTool/prompt.js' +import { getGoalContinuationPrompt } from '../services/goal/goalState.js' import { TICK_TAG } from './xml.js' import { logForDebugging } from '../utils/debug.js' import { loadMemoryPrompt } from '../memdir/memdir.js' @@ -505,6 +506,11 @@ ${CYBER_RISK_INSTRUCTION}`, ...(feature('KAIROS') || feature('KAIROS_BRIEF') ? [systemPromptSection('brief', () => getBriefSection())] : []), + DANGEROUS_uncachedSystemPromptSection( + 'goal_continuation', + () => getGoalContinuationPrompt(), + 'Goal state changes between turns', + ), ] const resolvedDynamicSections = diff --git a/src/query.ts b/src/query.ts index b2de65d4e2..c7b276bb64 100644 --- a/src/query.ts +++ b/src/query.ts @@ -5,6 +5,7 @@ import type { } from '@anthropic-ai/sdk/resources/index.mjs' import type { CanUseToolFn } from './hooks/useCanUseTool.js' import { FallbackTriggeredError } from './services/api/withRetry.js' +import { updateGoalTokens } from './services/goal/goalState.js' import { calculateTokenWarningState, estimateMaxTurnGrowth, @@ -1265,6 +1266,13 @@ async function* queryLoop( if (warningInfo) { yield createCacheWarningMessage(warningInfo) } + + // Update goal token usage + const totalTokens = + usage.input_tokens + + (usage.cache_creation_input_tokens ?? 0) + + (usage.cache_read_input_tokens ?? 0) + updateGoalTokens(totalTokens) } } diff --git a/src/services/goal/goalState.ts b/src/services/goal/goalState.ts new file mode 100644 index 0000000000..d6c3c375ec --- /dev/null +++ b/src/services/goal/goalState.ts @@ -0,0 +1,156 @@ +import { getSessionId } from '../../bootstrap/state.js' + +export type GoalStatus = 'active' | 'paused' | 'budget_limited' | 'complete' + +export type GoalState = { + objective: string + status: GoalStatus + tokenBudget: number | null + tokensUsed: number + startTime: number + pausedAt: number | null + accumulatedActiveMs: number +} + +const goals: Map = new Map() + +export function getGoal(sessionId?: string): GoalState | null { + return goals.get(sessionId ?? getSessionId()) ?? null +} + +export function setGoal( + objective: string, + tokenBudget?: number, + sessionId?: string, +): GoalState { + const validBudget = + tokenBudget !== undefined && + Number.isFinite(tokenBudget) && + tokenBudget >= 0 + ? tokenBudget + : null + const state: GoalState = { + objective, + status: 'active', + tokenBudget: validBudget, + tokensUsed: 0, + startTime: Date.now(), + pausedAt: null, + accumulatedActiveMs: 0, + } + goals.set(sessionId ?? getSessionId(), state) + return state +} + +export function clearGoal(sessionId?: string): void { + goals.delete(sessionId ?? getSessionId()) +} + +export function pauseGoal(sessionId?: string): boolean { + const goal = getGoal(sessionId) + if (!goal || goal.status !== 'active') return false + goal.accumulatedActiveMs += Date.now() - goal.startTime + goal.pausedAt = Date.now() + goal.status = 'paused' + return true +} + +export function resumeGoal(sessionId?: string): boolean { + const goal = getGoal(sessionId) + if (!goal || goal.status !== 'paused') return false + goal.pausedAt = null + goal.startTime = Date.now() + goal.status = 'active' + return true +} + +export function completeGoal(sessionId?: string): boolean { + const goal = getGoal(sessionId) + if (!goal) return false + goal.status = 'complete' + return true +} + +export function updateGoalTokens(usage: number, sessionId?: string): void { + const goal = getGoal(sessionId) + if (!goal || goal.status !== 'active') return + const validUsage = Number.isFinite(usage) && usage >= 0 ? usage : 0 + goal.tokensUsed += validUsage + if (goal.tokenBudget !== null && goal.tokensUsed >= goal.tokenBudget) { + goal.status = 'budget_limited' + } +} + +export function getActiveElapsedMs(goal: GoalState): number { + const ongoing = + goal.status === 'active' && goal.pausedAt === null + ? Date.now() - goal.startTime + : 0 + return goal.accumulatedActiveMs + ongoing +} + +export function getGoalContinuationPrompt(sessionId?: string): string | null { + const goal = getGoal(sessionId) + if (!goal || goal.status !== 'active') return null + + const elapsedSeconds = Math.floor(getActiveElapsedMs(goal) / 1000) + const budgetDisplay = + goal.tokenBudget !== null ? `${goal.tokenBudget}` : 'unlimited' + const remainingDisplay = + goal.tokenBudget !== null + ? `${Math.max(0, goal.tokenBudget - goal.tokensUsed)}` + : 'unlimited' + + return `Continue working toward the active goal. + + +${goal.objective} + + +Budget: +- Time spent: ${elapsedSeconds} seconds +- Tokens used: ${goal.tokensUsed} +- Token budget: ${budgetDisplay} +- Tokens remaining: ${remainingDisplay} + +Avoid repeating work that is already done. Choose the next concrete action toward the objective. + +Before deciding that the goal is achieved, perform a completion audit: +- Restate the objective as concrete deliverables or success criteria. +- Inspect relevant files, command output, test results, or other real evidence. +- Do not accept proxy signals as completion by themselves. +- Treat uncertainty as not achieved; do more verification or continue the work. +- Only mark the goal achieved when the objective has actually been achieved and no required work remains. + +If the objective is achieved, call the goal tool with action "complete" so usage accounting is preserved.` +} + +export function formatGoalStatus(sessionId?: string): string { + const goal = getGoal(sessionId) + if (!goal) return 'No active goal.' + + const elapsed = Math.floor(getActiveElapsedMs(goal) / 1000) + const minutes = Math.floor(elapsed / 60) + const seconds = elapsed % 60 + const timeStr = minutes > 0 ? `${minutes}m ${seconds}s` : `${seconds}s` + + const statusLabel: Record = { + active: 'Active', + paused: 'Paused', + budget_limited: 'Budget Limited', + complete: 'Complete', + } + + const lines = [ + `Goal: ${goal.objective}`, + `Status: ${statusLabel[goal.status]}`, + `Time: ${timeStr}`, + `Tokens: ${goal.tokensUsed}${goal.tokenBudget !== null ? ` / ${goal.tokenBudget}` : ''}`, + ] + + return lines.join('\n') +} + +export function clearAllGoals(): void { + goals.clear() +} diff --git a/src/tools.ts b/src/tools.ts index 08f26429be..68624d3729 100644 --- a/src/tools.ts +++ b/src/tools.ts @@ -87,6 +87,7 @@ import { EnterPlanModeTool } from '@claude-code-best/builtin-tools/tools/EnterPl import { EnterWorktreeTool } from '@claude-code-best/builtin-tools/tools/EnterWorktreeTool/EnterWorktreeTool.js' import { ExitWorktreeTool } from '@claude-code-best/builtin-tools/tools/ExitWorktreeTool/ExitWorktreeTool.js' import { ConfigTool } from '@claude-code-best/builtin-tools/tools/ConfigTool/ConfigTool.js' +import { GoalTool } from '@claude-code-best/builtin-tools/tools/GoalTool/GoalTool.js' import { LocalMemoryRecallTool } from '@claude-code-best/builtin-tools/tools/LocalMemoryRecallTool/LocalMemoryRecallTool.js' import { VaultHttpFetchTool } from '@claude-code-best/builtin-tools/tools/VaultHttpFetchTool/VaultHttpFetchTool.js' import { TaskCreateTool } from '@claude-code-best/builtin-tools/tools/TaskCreateTool/TaskCreateTool.js' @@ -261,6 +262,7 @@ export function getAllBaseTools(): Tools { ...(RemoteTriggerTool ? [RemoteTriggerTool] : []), ...(MonitorTool ? [MonitorTool] : []), BriefTool, + GoalTool, ...(SendUserFileTool ? [SendUserFileTool] : []), ...(PushNotificationTool ? [PushNotificationTool] : []), ...(SubscribePRTool ? [SubscribePRTool] : []),