chore: initialize recovered claude workspace
This commit is contained in:
523
src/services/PromptSuggestion/promptSuggestion.ts
Normal file
523
src/services/PromptSuggestion/promptSuggestion.ts
Normal file
@@ -0,0 +1,523 @@
|
||||
import { getIsNonInteractiveSession } from '../../bootstrap/state.js'
|
||||
import type { AppState } from '../../state/AppState.js'
|
||||
import type { Message } from '../../types/message.js'
|
||||
import { isAgentSwarmsEnabled } from '../../utils/agentSwarmsEnabled.js'
|
||||
import { count } from '../../utils/array.js'
|
||||
import { isEnvDefinedFalsy, isEnvTruthy } from '../../utils/envUtils.js'
|
||||
import { toError } from '../../utils/errors.js'
|
||||
import {
|
||||
type CacheSafeParams,
|
||||
createCacheSafeParams,
|
||||
runForkedAgent,
|
||||
} from '../../utils/forkedAgent.js'
|
||||
import type { REPLHookContext } from '../../utils/hooks/postSamplingHooks.js'
|
||||
import { logError } from '../../utils/log.js'
|
||||
import {
|
||||
createUserMessage,
|
||||
getLastAssistantMessage,
|
||||
} from '../../utils/messages.js'
|
||||
import { getInitialSettings } from '../../utils/settings/settings.js'
|
||||
import { isTeammate } from '../../utils/teammate.js'
|
||||
import { getFeatureValue_CACHED_MAY_BE_STALE } from '../analytics/growthbook.js'
|
||||
import {
|
||||
type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
logEvent,
|
||||
} from '../analytics/index.js'
|
||||
import { currentLimits } from '../claudeAiLimits.js'
|
||||
import { isSpeculationEnabled, startSpeculation } from './speculation.js'
|
||||
|
||||
let currentAbortController: AbortController | null = null
|
||||
|
||||
export type PromptVariant = 'user_intent' | 'stated_intent'
|
||||
|
||||
export function getPromptVariant(): PromptVariant {
|
||||
return 'user_intent'
|
||||
}
|
||||
|
||||
export function shouldEnablePromptSuggestion(): boolean {
|
||||
// Env var overrides everything (for testing)
|
||||
const envOverride = process.env.CLAUDE_CODE_ENABLE_PROMPT_SUGGESTION
|
||||
if (isEnvDefinedFalsy(envOverride)) {
|
||||
logEvent('tengu_prompt_suggestion_init', {
|
||||
enabled: false,
|
||||
source:
|
||||
'env' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
return false
|
||||
}
|
||||
if (isEnvTruthy(envOverride)) {
|
||||
logEvent('tengu_prompt_suggestion_init', {
|
||||
enabled: true,
|
||||
source:
|
||||
'env' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
return true
|
||||
}
|
||||
|
||||
// Keep default in sync with Config.tsx (settings toggle visibility)
|
||||
if (!getFeatureValue_CACHED_MAY_BE_STALE('tengu_chomp_inflection', false)) {
|
||||
logEvent('tengu_prompt_suggestion_init', {
|
||||
enabled: false,
|
||||
source:
|
||||
'growthbook' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
return false
|
||||
}
|
||||
|
||||
// Disable in non-interactive mode (print mode, piped input, SDK)
|
||||
if (getIsNonInteractiveSession()) {
|
||||
logEvent('tengu_prompt_suggestion_init', {
|
||||
enabled: false,
|
||||
source:
|
||||
'non_interactive' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
return false
|
||||
}
|
||||
|
||||
// Disable for swarm teammates (only leader should show suggestions)
|
||||
if (isAgentSwarmsEnabled() && isTeammate()) {
|
||||
logEvent('tengu_prompt_suggestion_init', {
|
||||
enabled: false,
|
||||
source:
|
||||
'swarm_teammate' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
return false
|
||||
}
|
||||
|
||||
const enabled = getInitialSettings()?.promptSuggestionEnabled !== false
|
||||
logEvent('tengu_prompt_suggestion_init', {
|
||||
enabled,
|
||||
source:
|
||||
'setting' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
return enabled
|
||||
}
|
||||
|
||||
export function abortPromptSuggestion(): void {
|
||||
if (currentAbortController) {
|
||||
currentAbortController.abort()
|
||||
currentAbortController = null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a suppression reason if suggestions should not be generated,
|
||||
* or null if generation is allowed. Shared by main and pipelined paths.
|
||||
*/
|
||||
export function getSuggestionSuppressReason(appState: AppState): string | null {
|
||||
if (!appState.promptSuggestionEnabled) return 'disabled'
|
||||
if (appState.pendingWorkerRequest || appState.pendingSandboxRequest)
|
||||
return 'pending_permission'
|
||||
if (appState.elicitation.queue.length > 0) return 'elicitation_active'
|
||||
if (appState.toolPermissionContext.mode === 'plan') return 'plan_mode'
|
||||
if (
|
||||
process.env.USER_TYPE === 'external' &&
|
||||
currentLimits.status !== 'allowed'
|
||||
)
|
||||
return 'rate_limit'
|
||||
return null
|
||||
}
|
||||
|
||||
/**
|
||||
* Shared guard + generation logic used by both CLI TUI and SDK push paths.
|
||||
* Returns the suggestion with metadata, or null if suppressed/filtered.
|
||||
*/
|
||||
export async function tryGenerateSuggestion(
|
||||
abortController: AbortController,
|
||||
messages: Message[],
|
||||
getAppState: () => AppState,
|
||||
cacheSafeParams: CacheSafeParams,
|
||||
source?: 'cli' | 'sdk',
|
||||
): Promise<{
|
||||
suggestion: string
|
||||
promptId: PromptVariant
|
||||
generationRequestId: string | null
|
||||
} | null> {
|
||||
if (abortController.signal.aborted) {
|
||||
logSuggestionSuppressed('aborted', undefined, undefined, source)
|
||||
return null
|
||||
}
|
||||
|
||||
const assistantTurnCount = count(messages, m => m.type === 'assistant')
|
||||
if (assistantTurnCount < 2) {
|
||||
logSuggestionSuppressed('early_conversation', undefined, undefined, source)
|
||||
return null
|
||||
}
|
||||
|
||||
const lastAssistantMessage = getLastAssistantMessage(messages)
|
||||
if (lastAssistantMessage?.isApiErrorMessage) {
|
||||
logSuggestionSuppressed('last_response_error', undefined, undefined, source)
|
||||
return null
|
||||
}
|
||||
const cacheReason = getParentCacheSuppressReason(lastAssistantMessage)
|
||||
if (cacheReason) {
|
||||
logSuggestionSuppressed(cacheReason, undefined, undefined, source)
|
||||
return null
|
||||
}
|
||||
|
||||
const appState = getAppState()
|
||||
const suppressReason = getSuggestionSuppressReason(appState)
|
||||
if (suppressReason) {
|
||||
logSuggestionSuppressed(suppressReason, undefined, undefined, source)
|
||||
return null
|
||||
}
|
||||
|
||||
const promptId = getPromptVariant()
|
||||
const { suggestion, generationRequestId } = await generateSuggestion(
|
||||
abortController,
|
||||
promptId,
|
||||
cacheSafeParams,
|
||||
)
|
||||
if (abortController.signal.aborted) {
|
||||
logSuggestionSuppressed('aborted', undefined, undefined, source)
|
||||
return null
|
||||
}
|
||||
if (!suggestion) {
|
||||
logSuggestionSuppressed('empty', undefined, promptId, source)
|
||||
return null
|
||||
}
|
||||
if (shouldFilterSuggestion(suggestion, promptId, source)) return null
|
||||
|
||||
return { suggestion, promptId, generationRequestId }
|
||||
}
|
||||
|
||||
export async function executePromptSuggestion(
|
||||
context: REPLHookContext,
|
||||
): Promise<void> {
|
||||
if (context.querySource !== 'repl_main_thread') return
|
||||
|
||||
currentAbortController = new AbortController()
|
||||
const abortController = currentAbortController
|
||||
const cacheSafeParams = createCacheSafeParams(context)
|
||||
|
||||
try {
|
||||
const result = await tryGenerateSuggestion(
|
||||
abortController,
|
||||
context.messages,
|
||||
context.toolUseContext.getAppState,
|
||||
cacheSafeParams,
|
||||
'cli',
|
||||
)
|
||||
if (!result) return
|
||||
|
||||
context.toolUseContext.setAppState(prev => ({
|
||||
...prev,
|
||||
promptSuggestion: {
|
||||
text: result.suggestion,
|
||||
promptId: result.promptId,
|
||||
shownAt: 0,
|
||||
acceptedAt: 0,
|
||||
generationRequestId: result.generationRequestId,
|
||||
},
|
||||
}))
|
||||
|
||||
if (isSpeculationEnabled() && result.suggestion) {
|
||||
void startSpeculation(
|
||||
result.suggestion,
|
||||
context,
|
||||
context.toolUseContext.setAppState,
|
||||
false,
|
||||
cacheSafeParams,
|
||||
)
|
||||
}
|
||||
} catch (error) {
|
||||
if (
|
||||
error instanceof Error &&
|
||||
(error.name === 'AbortError' || error.name === 'APIUserAbortError')
|
||||
) {
|
||||
logSuggestionSuppressed('aborted', undefined, undefined, 'cli')
|
||||
return
|
||||
}
|
||||
logError(toError(error))
|
||||
} finally {
|
||||
if (currentAbortController === abortController) {
|
||||
currentAbortController = null
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const MAX_PARENT_UNCACHED_TOKENS = 10_000
|
||||
|
||||
export function getParentCacheSuppressReason(
|
||||
lastAssistantMessage: ReturnType<typeof getLastAssistantMessage>,
|
||||
): string | null {
|
||||
if (!lastAssistantMessage) return null
|
||||
|
||||
const usage = lastAssistantMessage.message.usage
|
||||
const inputTokens = usage.input_tokens ?? 0
|
||||
const cacheWriteTokens = usage.cache_creation_input_tokens ?? 0
|
||||
// The fork re-processes the parent's output (never cached) plus its own prompt.
|
||||
const outputTokens = usage.output_tokens ?? 0
|
||||
|
||||
return inputTokens + cacheWriteTokens + outputTokens >
|
||||
MAX_PARENT_UNCACHED_TOKENS
|
||||
? 'cache_cold'
|
||||
: null
|
||||
}
|
||||
|
||||
const SUGGESTION_PROMPT = `[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]
|
||||
|
||||
FIRST: Look at the user's recent messages and original request.
|
||||
|
||||
Your job is to predict what THEY would type - not what you think they should do.
|
||||
|
||||
THE TEST: Would they think "I was just about to type that"?
|
||||
|
||||
EXAMPLES:
|
||||
User asked "fix the bug and run tests", bug is fixed → "run the tests"
|
||||
After code written → "try it out"
|
||||
Claude offers options → suggest the one the user would likely pick, based on conversation
|
||||
Claude asks to continue → "yes" or "go ahead"
|
||||
Task complete, obvious follow-up → "commit this" or "push it"
|
||||
After error or misunderstanding → silence (let them assess/correct)
|
||||
|
||||
Be specific: "run the tests" beats "continue".
|
||||
|
||||
NEVER SUGGEST:
|
||||
- Evaluative ("looks good", "thanks")
|
||||
- Questions ("what about...?")
|
||||
- Claude-voice ("Let me...", "I'll...", "Here's...")
|
||||
- New ideas they didn't ask about
|
||||
- Multiple sentences
|
||||
|
||||
Stay silent if the next step isn't obvious from what the user said.
|
||||
|
||||
Format: 2-12 words, match the user's style. Or nothing.
|
||||
|
||||
Reply with ONLY the suggestion, no quotes or explanation.`
|
||||
|
||||
const SUGGESTION_PROMPTS: Record<PromptVariant, string> = {
|
||||
user_intent: SUGGESTION_PROMPT,
|
||||
stated_intent: SUGGESTION_PROMPT,
|
||||
}
|
||||
|
||||
export async function generateSuggestion(
|
||||
abortController: AbortController,
|
||||
promptId: PromptVariant,
|
||||
cacheSafeParams: CacheSafeParams,
|
||||
): Promise<{ suggestion: string | null; generationRequestId: string | null }> {
|
||||
const prompt = SUGGESTION_PROMPTS[promptId]
|
||||
|
||||
// Deny tools via callback, NOT by passing tools:[] - that busts cache (0% hit)
|
||||
const canUseTool = async () => ({
|
||||
behavior: 'deny' as const,
|
||||
message: 'No tools needed for suggestion',
|
||||
decisionReason: { type: 'other' as const, reason: 'suggestion only' },
|
||||
})
|
||||
|
||||
// DO NOT override any API parameter that differs from the parent request.
|
||||
// The fork piggybacks on the main thread's prompt cache by sending identical
|
||||
// cache-key params. The billing cache key includes more than just
|
||||
// system/tools/model/messages/thinking — empirically, setting effortValue
|
||||
// or maxOutputTokens on the fork (even via output_config or getAppState)
|
||||
// busts cache. PR #18143 tried effort:'low' and caused a 45x spike in cache
|
||||
// writes (92.7% → 61% hit rate). The only safe overrides are:
|
||||
// - abortController (not sent to API)
|
||||
// - skipTranscript (client-side only)
|
||||
// - skipCacheWrite (controls cache_control markers, not the cache key)
|
||||
// - canUseTool (client-side permission check)
|
||||
const result = await runForkedAgent({
|
||||
promptMessages: [createUserMessage({ content: prompt })],
|
||||
cacheSafeParams, // Don't override tools/thinking settings - busts cache
|
||||
canUseTool,
|
||||
querySource: 'prompt_suggestion',
|
||||
forkLabel: 'prompt_suggestion',
|
||||
overrides: {
|
||||
abortController,
|
||||
},
|
||||
skipTranscript: true,
|
||||
skipCacheWrite: true,
|
||||
})
|
||||
|
||||
// Check ALL messages - model may loop (try tool → denied → text in next message)
|
||||
// Also extract the requestId from the first assistant message for RL dataset joins
|
||||
const firstAssistantMsg = result.messages.find(m => m.type === 'assistant')
|
||||
const generationRequestId =
|
||||
firstAssistantMsg?.type === 'assistant'
|
||||
? (firstAssistantMsg.requestId ?? null)
|
||||
: null
|
||||
|
||||
for (const msg of result.messages) {
|
||||
if (msg.type !== 'assistant') continue
|
||||
const textBlock = msg.message.content.find(b => b.type === 'text')
|
||||
if (textBlock?.type === 'text') {
|
||||
const suggestion = textBlock.text.trim()
|
||||
if (suggestion) {
|
||||
return { suggestion, generationRequestId }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { suggestion: null, generationRequestId }
|
||||
}
|
||||
|
||||
export function shouldFilterSuggestion(
|
||||
suggestion: string | null,
|
||||
promptId: PromptVariant,
|
||||
source?: 'cli' | 'sdk',
|
||||
): boolean {
|
||||
if (!suggestion) {
|
||||
logSuggestionSuppressed('empty', undefined, promptId, source)
|
||||
return true
|
||||
}
|
||||
|
||||
const lower = suggestion.toLowerCase()
|
||||
const wordCount = suggestion.trim().split(/\s+/).length
|
||||
|
||||
const filters: Array<[string, () => boolean]> = [
|
||||
['done', () => lower === 'done'],
|
||||
[
|
||||
'meta_text',
|
||||
() =>
|
||||
lower === 'nothing found' ||
|
||||
lower === 'nothing found.' ||
|
||||
lower.startsWith('nothing to suggest') ||
|
||||
lower.startsWith('no suggestion') ||
|
||||
// Model spells out the prompt's "stay silent" instruction
|
||||
/\bsilence is\b|\bstay(s|ing)? silent\b/.test(lower) ||
|
||||
// Model outputs bare "silence" wrapped in punctuation/whitespace
|
||||
/^\W*silence\W*$/.test(lower),
|
||||
],
|
||||
[
|
||||
'meta_wrapped',
|
||||
// Model wraps meta-reasoning in parens/brackets: (silence — ...), [no suggestion]
|
||||
() => /^\(.*\)$|^\[.*\]$/.test(suggestion),
|
||||
],
|
||||
[
|
||||
'error_message',
|
||||
() =>
|
||||
lower.startsWith('api error:') ||
|
||||
lower.startsWith('prompt is too long') ||
|
||||
lower.startsWith('request timed out') ||
|
||||
lower.startsWith('invalid api key') ||
|
||||
lower.startsWith('image was too large'),
|
||||
],
|
||||
['prefixed_label', () => /^\w+:\s/.test(suggestion)],
|
||||
[
|
||||
'too_few_words',
|
||||
() => {
|
||||
if (wordCount >= 2) return false
|
||||
// Allow slash commands — these are valid user commands
|
||||
if (suggestion.startsWith('/')) return false
|
||||
// Allow common single-word inputs that are valid user commands
|
||||
const ALLOWED_SINGLE_WORDS = new Set([
|
||||
// Affirmatives
|
||||
'yes',
|
||||
'yeah',
|
||||
'yep',
|
||||
'yea',
|
||||
'yup',
|
||||
'sure',
|
||||
'ok',
|
||||
'okay',
|
||||
// Actions
|
||||
'push',
|
||||
'commit',
|
||||
'deploy',
|
||||
'stop',
|
||||
'continue',
|
||||
'check',
|
||||
'exit',
|
||||
'quit',
|
||||
// Negation
|
||||
'no',
|
||||
])
|
||||
return !ALLOWED_SINGLE_WORDS.has(lower)
|
||||
},
|
||||
],
|
||||
['too_many_words', () => wordCount > 12],
|
||||
['too_long', () => suggestion.length >= 100],
|
||||
['multiple_sentences', () => /[.!?]\s+[A-Z]/.test(suggestion)],
|
||||
['has_formatting', () => /[\n*]|\*\*/.test(suggestion)],
|
||||
[
|
||||
'evaluative',
|
||||
() =>
|
||||
/thanks|thank you|looks good|sounds good|that works|that worked|that's all|nice|great|perfect|makes sense|awesome|excellent/.test(
|
||||
lower,
|
||||
),
|
||||
],
|
||||
[
|
||||
'claude_voice',
|
||||
() =>
|
||||
/^(let me|i'll|i've|i'm|i can|i would|i think|i notice|here's|here is|here are|that's|this is|this will|you can|you should|you could|sure,|of course|certainly)/i.test(
|
||||
suggestion,
|
||||
),
|
||||
],
|
||||
]
|
||||
|
||||
for (const [reason, check] of filters) {
|
||||
if (check()) {
|
||||
logSuggestionSuppressed(reason, suggestion, promptId, source)
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
/**
|
||||
* Log acceptance/ignoring of a prompt suggestion. Used by the SDK push path
|
||||
* to track outcomes when the next user message arrives.
|
||||
*/
|
||||
export function logSuggestionOutcome(
|
||||
suggestion: string,
|
||||
userInput: string,
|
||||
emittedAt: number,
|
||||
promptId: PromptVariant,
|
||||
generationRequestId: string | null,
|
||||
): void {
|
||||
const similarity =
|
||||
Math.round((userInput.length / (suggestion.length || 1)) * 100) / 100
|
||||
const wasAccepted = userInput === suggestion
|
||||
const timeMs = Math.max(0, Date.now() - emittedAt)
|
||||
|
||||
logEvent('tengu_prompt_suggestion', {
|
||||
source: 'sdk' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
outcome: (wasAccepted
|
||||
? 'accepted'
|
||||
: 'ignored') as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
prompt_id:
|
||||
promptId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
...(generationRequestId && {
|
||||
generationRequestId:
|
||||
generationRequestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}),
|
||||
...(wasAccepted && {
|
||||
timeToAcceptMs: timeMs,
|
||||
}),
|
||||
...(!wasAccepted && { timeToIgnoreMs: timeMs }),
|
||||
similarity,
|
||||
...(process.env.USER_TYPE === 'ant' && {
|
||||
suggestion:
|
||||
suggestion as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
userInput:
|
||||
userInput as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
export function logSuggestionSuppressed(
|
||||
reason: string,
|
||||
suggestion?: string,
|
||||
promptId?: PromptVariant,
|
||||
source?: 'cli' | 'sdk',
|
||||
): void {
|
||||
const resolvedPromptId = promptId ?? getPromptVariant()
|
||||
logEvent('tengu_prompt_suggestion', {
|
||||
...(source && {
|
||||
source:
|
||||
source as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}),
|
||||
outcome:
|
||||
'suppressed' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
reason:
|
||||
reason as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
prompt_id:
|
||||
resolvedPromptId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
...(process.env.USER_TYPE === 'ant' &&
|
||||
suggestion && {
|
||||
suggestion:
|
||||
suggestion as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}),
|
||||
})
|
||||
}
|
||||
991
src/services/PromptSuggestion/speculation.ts
Normal file
991
src/services/PromptSuggestion/speculation.ts
Normal file
@@ -0,0 +1,991 @@
|
||||
import { randomUUID } from 'crypto'
|
||||
import { rm } from 'fs'
|
||||
import { appendFile, copyFile, mkdir } from 'fs/promises'
|
||||
import { dirname, isAbsolute, join, relative } from 'path'
|
||||
import { getCwdState } from '../../bootstrap/state.js'
|
||||
import type { CompletionBoundary } from '../../state/AppStateStore.js'
|
||||
import {
|
||||
type AppState,
|
||||
IDLE_SPECULATION_STATE,
|
||||
type SpeculationResult,
|
||||
type SpeculationState,
|
||||
} from '../../state/AppStateStore.js'
|
||||
import { commandHasAnyCd } from '../../tools/BashTool/bashPermissions.js'
|
||||
import { checkReadOnlyConstraints } from '../../tools/BashTool/readOnlyValidation.js'
|
||||
import type { SpeculationAcceptMessage } from '../../types/logs.js'
|
||||
import type { Message } from '../../types/message.js'
|
||||
import { createChildAbortController } from '../../utils/abortController.js'
|
||||
import { count } from '../../utils/array.js'
|
||||
import { getGlobalConfig } from '../../utils/config.js'
|
||||
import { logForDebugging } from '../../utils/debug.js'
|
||||
import { errorMessage } from '../../utils/errors.js'
|
||||
import {
|
||||
type FileStateCache,
|
||||
mergeFileStateCaches,
|
||||
READ_FILE_STATE_CACHE_SIZE,
|
||||
} from '../../utils/fileStateCache.js'
|
||||
import {
|
||||
type CacheSafeParams,
|
||||
createCacheSafeParams,
|
||||
runForkedAgent,
|
||||
} from '../../utils/forkedAgent.js'
|
||||
import { formatDuration, formatNumber } from '../../utils/format.js'
|
||||
import type { REPLHookContext } from '../../utils/hooks/postSamplingHooks.js'
|
||||
import { logError } from '../../utils/log.js'
|
||||
import type { SetAppState } from '../../utils/messageQueueManager.js'
|
||||
import {
|
||||
createSystemMessage,
|
||||
createUserMessage,
|
||||
INTERRUPT_MESSAGE,
|
||||
INTERRUPT_MESSAGE_FOR_TOOL_USE,
|
||||
} from '../../utils/messages.js'
|
||||
import { getClaudeTempDir } from '../../utils/permissions/filesystem.js'
|
||||
import { extractReadFilesFromMessages } from '../../utils/queryHelpers.js'
|
||||
import { getTranscriptPath } from '../../utils/sessionStorage.js'
|
||||
import { jsonStringify } from '../../utils/slowOperations.js'
|
||||
import {
|
||||
type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
logEvent,
|
||||
} from '../analytics/index.js'
|
||||
import {
|
||||
generateSuggestion,
|
||||
getPromptVariant,
|
||||
getSuggestionSuppressReason,
|
||||
logSuggestionSuppressed,
|
||||
shouldFilterSuggestion,
|
||||
} from './promptSuggestion.js'
|
||||
|
||||
const MAX_SPECULATION_TURNS = 20
|
||||
const MAX_SPECULATION_MESSAGES = 100
|
||||
|
||||
const WRITE_TOOLS = new Set(['Edit', 'Write', 'NotebookEdit'])
|
||||
const SAFE_READ_ONLY_TOOLS = new Set([
|
||||
'Read',
|
||||
'Glob',
|
||||
'Grep',
|
||||
'ToolSearch',
|
||||
'LSP',
|
||||
'TaskGet',
|
||||
'TaskList',
|
||||
])
|
||||
|
||||
function safeRemoveOverlay(overlayPath: string): void {
|
||||
rm(
|
||||
overlayPath,
|
||||
{ recursive: true, force: true, maxRetries: 3, retryDelay: 100 },
|
||||
() => {},
|
||||
)
|
||||
}
|
||||
|
||||
function getOverlayPath(id: string): string {
|
||||
return join(getClaudeTempDir(), 'speculation', String(process.pid), id)
|
||||
}
|
||||
|
||||
function denySpeculation(
|
||||
message: string,
|
||||
reason: string,
|
||||
): {
|
||||
behavior: 'deny'
|
||||
message: string
|
||||
decisionReason: { type: 'other'; reason: string }
|
||||
} {
|
||||
return {
|
||||
behavior: 'deny',
|
||||
message,
|
||||
decisionReason: { type: 'other', reason },
|
||||
}
|
||||
}
|
||||
|
||||
async function copyOverlayToMain(
|
||||
overlayPath: string,
|
||||
writtenPaths: Set<string>,
|
||||
cwd: string,
|
||||
): Promise<boolean> {
|
||||
let allCopied = true
|
||||
for (const rel of writtenPaths) {
|
||||
const src = join(overlayPath, rel)
|
||||
const dest = join(cwd, rel)
|
||||
try {
|
||||
await mkdir(dirname(dest), { recursive: true })
|
||||
await copyFile(src, dest)
|
||||
} catch {
|
||||
allCopied = false
|
||||
logForDebugging(`[Speculation] Failed to copy ${rel} to main`)
|
||||
}
|
||||
}
|
||||
return allCopied
|
||||
}
|
||||
|
||||
export type ActiveSpeculationState = Extract<
|
||||
SpeculationState,
|
||||
{ status: 'active' }
|
||||
>
|
||||
|
||||
function logSpeculation(
|
||||
id: string,
|
||||
outcome: 'accepted' | 'aborted' | 'error',
|
||||
startTime: number,
|
||||
suggestionLength: number,
|
||||
messages: Message[],
|
||||
boundary: CompletionBoundary | null,
|
||||
extras?: Record<string, string | number | boolean | undefined>,
|
||||
): void {
|
||||
logEvent('tengu_speculation', {
|
||||
speculation_id:
|
||||
id as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
outcome:
|
||||
outcome as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
duration_ms: Date.now() - startTime,
|
||||
suggestion_length: suggestionLength,
|
||||
tools_executed: countToolsInMessages(messages),
|
||||
completed: boundary !== null,
|
||||
boundary_type: boundary?.type as
|
||||
| AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
|
||||
| undefined,
|
||||
boundary_tool: getBoundaryTool(boundary) as
|
||||
| AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
|
||||
| undefined,
|
||||
boundary_detail: getBoundaryDetail(boundary) as
|
||||
| AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
|
||||
| undefined,
|
||||
...extras,
|
||||
})
|
||||
}
|
||||
|
||||
function countToolsInMessages(messages: Message[]): number {
|
||||
const blocks = messages
|
||||
.filter(isUserMessageWithArrayContent)
|
||||
.flatMap(m => m.message.content)
|
||||
.filter(
|
||||
(b): b is { type: string; is_error?: boolean } =>
|
||||
typeof b === 'object' && b !== null && 'type' in b,
|
||||
)
|
||||
return count(blocks, b => b.type === 'tool_result' && !b.is_error)
|
||||
}
|
||||
|
||||
function getBoundaryTool(
|
||||
boundary: CompletionBoundary | null,
|
||||
): string | undefined {
|
||||
if (!boundary) return undefined
|
||||
switch (boundary.type) {
|
||||
case 'bash':
|
||||
return 'Bash'
|
||||
case 'edit':
|
||||
case 'denied_tool':
|
||||
return boundary.toolName
|
||||
case 'complete':
|
||||
return undefined
|
||||
}
|
||||
}
|
||||
|
||||
function getBoundaryDetail(
|
||||
boundary: CompletionBoundary | null,
|
||||
): string | undefined {
|
||||
if (!boundary) return undefined
|
||||
switch (boundary.type) {
|
||||
case 'bash':
|
||||
return boundary.command.slice(0, 200)
|
||||
case 'edit':
|
||||
return boundary.filePath
|
||||
case 'denied_tool':
|
||||
return boundary.detail
|
||||
case 'complete':
|
||||
return undefined
|
||||
}
|
||||
}
|
||||
|
||||
function isUserMessageWithArrayContent(
|
||||
m: Message,
|
||||
): m is Message & { message: { content: unknown[] } } {
|
||||
return m.type === 'user' && 'message' in m && Array.isArray(m.message.content)
|
||||
}
|
||||
|
||||
export function prepareMessagesForInjection(messages: Message[]): Message[] {
|
||||
// Find tool_use IDs that have SUCCESSFUL results (not errors/interruptions)
|
||||
// Pending tool_use blocks (no result) and interrupted ones will be stripped
|
||||
type ToolResult = {
|
||||
type: 'tool_result'
|
||||
tool_use_id: string
|
||||
is_error?: boolean
|
||||
content?: unknown
|
||||
}
|
||||
const isToolResult = (b: unknown): b is ToolResult =>
|
||||
typeof b === 'object' &&
|
||||
b !== null &&
|
||||
(b as ToolResult).type === 'tool_result' &&
|
||||
typeof (b as ToolResult).tool_use_id === 'string'
|
||||
const isSuccessful = (b: ToolResult) =>
|
||||
!b.is_error &&
|
||||
!(
|
||||
typeof b.content === 'string' &&
|
||||
b.content.includes(INTERRUPT_MESSAGE_FOR_TOOL_USE)
|
||||
)
|
||||
|
||||
const toolIdsWithSuccessfulResults = new Set(
|
||||
messages
|
||||
.filter(isUserMessageWithArrayContent)
|
||||
.flatMap(m => m.message.content)
|
||||
.filter(isToolResult)
|
||||
.filter(isSuccessful)
|
||||
.map(b => b.tool_use_id),
|
||||
)
|
||||
|
||||
const keep = (b: {
|
||||
type: string
|
||||
id?: string
|
||||
tool_use_id?: string
|
||||
text?: string
|
||||
}) =>
|
||||
b.type !== 'thinking' &&
|
||||
b.type !== 'redacted_thinking' &&
|
||||
!(b.type === 'tool_use' && !toolIdsWithSuccessfulResults.has(b.id!)) &&
|
||||
!(
|
||||
b.type === 'tool_result' &&
|
||||
!toolIdsWithSuccessfulResults.has(b.tool_use_id!)
|
||||
) &&
|
||||
// Abort during speculation yields a standalone interrupt user message
|
||||
// (query.ts createUserInterruptionMessage). Strip it so it isn't surfaced
|
||||
// to the model as real user input.
|
||||
!(
|
||||
b.type === 'text' &&
|
||||
(b.text === INTERRUPT_MESSAGE ||
|
||||
b.text === INTERRUPT_MESSAGE_FOR_TOOL_USE)
|
||||
)
|
||||
|
||||
return messages
|
||||
.map(msg => {
|
||||
if (!('message' in msg) || !Array.isArray(msg.message.content)) return msg
|
||||
const content = msg.message.content.filter(keep)
|
||||
if (content.length === msg.message.content.length) return msg
|
||||
if (content.length === 0) return null
|
||||
// Drop messages where all remaining blocks are whitespace-only text
|
||||
// (API rejects these with 400: "text content blocks must contain non-whitespace text")
|
||||
const hasNonWhitespaceContent = content.some(
|
||||
(b: { type: string; text?: string }) =>
|
||||
b.type !== 'text' || (b.text !== undefined && b.text.trim() !== ''),
|
||||
)
|
||||
if (!hasNonWhitespaceContent) return null
|
||||
return { ...msg, message: { ...msg.message, content } } as typeof msg
|
||||
})
|
||||
.filter((m): m is Message => m !== null)
|
||||
}
|
||||
|
||||
function createSpeculationFeedbackMessage(
|
||||
messages: Message[],
|
||||
boundary: CompletionBoundary | null,
|
||||
timeSavedMs: number,
|
||||
sessionTotalMs: number,
|
||||
): Message | null {
|
||||
if (process.env.USER_TYPE !== 'ant') return null
|
||||
|
||||
if (messages.length === 0 || timeSavedMs === 0) return null
|
||||
|
||||
const toolUses = countToolsInMessages(messages)
|
||||
const tokens = boundary?.type === 'complete' ? boundary.outputTokens : null
|
||||
|
||||
const parts = []
|
||||
if (toolUses > 0) {
|
||||
parts.push(`Speculated ${toolUses} tool ${toolUses === 1 ? 'use' : 'uses'}`)
|
||||
} else {
|
||||
const turns = messages.length
|
||||
parts.push(`Speculated ${turns} ${turns === 1 ? 'turn' : 'turns'}`)
|
||||
}
|
||||
|
||||
if (tokens !== null) {
|
||||
parts.push(`${formatNumber(tokens)} tokens`)
|
||||
}
|
||||
|
||||
const savedText = `+${formatDuration(timeSavedMs)} saved`
|
||||
const sessionSuffix =
|
||||
sessionTotalMs !== timeSavedMs
|
||||
? ` (${formatDuration(sessionTotalMs)} this session)`
|
||||
: ''
|
||||
|
||||
return createSystemMessage(
|
||||
`[ANT-ONLY] ${parts.join(' · ')} · ${savedText}${sessionSuffix}`,
|
||||
'warning',
|
||||
)
|
||||
}
|
||||
|
||||
function updateActiveSpeculationState(
|
||||
setAppState: SetAppState,
|
||||
updater: (state: ActiveSpeculationState) => Partial<ActiveSpeculationState>,
|
||||
): void {
|
||||
setAppState(prev => {
|
||||
if (prev.speculation.status !== 'active') return prev
|
||||
const current = prev.speculation as ActiveSpeculationState
|
||||
const updates = updater(current)
|
||||
// Check if any values actually changed to avoid unnecessary re-renders
|
||||
const hasChanges = Object.entries(updates).some(
|
||||
([key, value]) => current[key as keyof ActiveSpeculationState] !== value,
|
||||
)
|
||||
if (!hasChanges) return prev
|
||||
return {
|
||||
...prev,
|
||||
speculation: { ...current, ...updates },
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
function resetSpeculationState(setAppState: SetAppState): void {
|
||||
setAppState(prev => {
|
||||
if (prev.speculation.status === 'idle') return prev
|
||||
return { ...prev, speculation: IDLE_SPECULATION_STATE }
|
||||
})
|
||||
}
|
||||
|
||||
export function isSpeculationEnabled(): boolean {
|
||||
const enabled =
|
||||
process.env.USER_TYPE === 'ant' &&
|
||||
(getGlobalConfig().speculationEnabled ?? true)
|
||||
logForDebugging(`[Speculation] enabled=${enabled}`)
|
||||
return enabled
|
||||
}
|
||||
|
||||
async function generatePipelinedSuggestion(
|
||||
context: REPLHookContext,
|
||||
suggestionText: string,
|
||||
speculatedMessages: Message[],
|
||||
setAppState: SetAppState,
|
||||
parentAbortController: AbortController,
|
||||
): Promise<void> {
|
||||
try {
|
||||
const appState = context.toolUseContext.getAppState()
|
||||
const suppressReason = getSuggestionSuppressReason(appState)
|
||||
if (suppressReason) {
|
||||
logSuggestionSuppressed(`pipeline_${suppressReason}`)
|
||||
return
|
||||
}
|
||||
|
||||
const augmentedContext: REPLHookContext = {
|
||||
...context,
|
||||
messages: [
|
||||
...context.messages,
|
||||
createUserMessage({ content: suggestionText }),
|
||||
...speculatedMessages,
|
||||
],
|
||||
}
|
||||
|
||||
const pipelineAbortController = createChildAbortController(
|
||||
parentAbortController,
|
||||
)
|
||||
if (pipelineAbortController.signal.aborted) return
|
||||
|
||||
const promptId = getPromptVariant()
|
||||
const { suggestion, generationRequestId } = await generateSuggestion(
|
||||
pipelineAbortController,
|
||||
promptId,
|
||||
createCacheSafeParams(augmentedContext),
|
||||
)
|
||||
|
||||
if (pipelineAbortController.signal.aborted) return
|
||||
if (shouldFilterSuggestion(suggestion, promptId)) return
|
||||
|
||||
logForDebugging(
|
||||
`[Speculation] Pipelined suggestion: "${suggestion!.slice(0, 50)}..."`,
|
||||
)
|
||||
updateActiveSpeculationState(setAppState, () => ({
|
||||
pipelinedSuggestion: {
|
||||
text: suggestion!,
|
||||
promptId,
|
||||
generationRequestId,
|
||||
},
|
||||
}))
|
||||
} catch (error) {
|
||||
if (error instanceof Error && error.name === 'AbortError') return
|
||||
logForDebugging(
|
||||
`[Speculation] Pipelined suggestion failed: ${errorMessage(error)}`,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
export async function startSpeculation(
|
||||
suggestionText: string,
|
||||
context: REPLHookContext,
|
||||
setAppState: (f: (prev: AppState) => AppState) => void,
|
||||
isPipelined = false,
|
||||
cacheSafeParams?: CacheSafeParams,
|
||||
): Promise<void> {
|
||||
if (!isSpeculationEnabled()) return
|
||||
|
||||
// Abort any existing speculation before starting a new one
|
||||
abortSpeculation(setAppState)
|
||||
|
||||
const id = randomUUID().slice(0, 8)
|
||||
|
||||
const abortController = createChildAbortController(
|
||||
context.toolUseContext.abortController,
|
||||
)
|
||||
|
||||
if (abortController.signal.aborted) return
|
||||
|
||||
const startTime = Date.now()
|
||||
const messagesRef = { current: [] as Message[] }
|
||||
const writtenPathsRef = { current: new Set<string>() }
|
||||
const overlayPath = getOverlayPath(id)
|
||||
const cwd = getCwdState()
|
||||
|
||||
try {
|
||||
await mkdir(overlayPath, { recursive: true })
|
||||
} catch {
|
||||
logForDebugging('[Speculation] Failed to create overlay directory')
|
||||
return
|
||||
}
|
||||
|
||||
const contextRef = { current: context }
|
||||
|
||||
setAppState(prev => ({
|
||||
...prev,
|
||||
speculation: {
|
||||
status: 'active',
|
||||
id,
|
||||
abort: () => abortController.abort(),
|
||||
startTime,
|
||||
messagesRef,
|
||||
writtenPathsRef,
|
||||
boundary: null,
|
||||
suggestionLength: suggestionText.length,
|
||||
toolUseCount: 0,
|
||||
isPipelined,
|
||||
contextRef,
|
||||
},
|
||||
}))
|
||||
|
||||
logForDebugging(`[Speculation] Starting speculation ${id}`)
|
||||
|
||||
try {
|
||||
const result = await runForkedAgent({
|
||||
promptMessages: [createUserMessage({ content: suggestionText })],
|
||||
cacheSafeParams: cacheSafeParams ?? createCacheSafeParams(context),
|
||||
skipTranscript: true,
|
||||
canUseTool: async (tool, input) => {
|
||||
const isWriteTool = WRITE_TOOLS.has(tool.name)
|
||||
const isSafeReadOnlyTool = SAFE_READ_ONLY_TOOLS.has(tool.name)
|
||||
|
||||
// Check permission mode BEFORE allowing file edits
|
||||
if (isWriteTool) {
|
||||
const appState = context.toolUseContext.getAppState()
|
||||
const { mode, isBypassPermissionsModeAvailable } =
|
||||
appState.toolPermissionContext
|
||||
|
||||
const canAutoAcceptEdits =
|
||||
mode === 'acceptEdits' ||
|
||||
mode === 'bypassPermissions' ||
|
||||
(mode === 'plan' && isBypassPermissionsModeAvailable)
|
||||
|
||||
if (!canAutoAcceptEdits) {
|
||||
logForDebugging(`[Speculation] Stopping at file edit: ${tool.name}`)
|
||||
const editPath = (
|
||||
'file_path' in input ? input.file_path : undefined
|
||||
) as string | undefined
|
||||
updateActiveSpeculationState(setAppState, () => ({
|
||||
boundary: {
|
||||
type: 'edit',
|
||||
toolName: tool.name,
|
||||
filePath: editPath ?? '',
|
||||
completedAt: Date.now(),
|
||||
},
|
||||
}))
|
||||
abortController.abort()
|
||||
return denySpeculation(
|
||||
'Speculation paused: file edit requires permission',
|
||||
'speculation_edit_boundary',
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// Handle file path rewriting for overlay isolation
|
||||
if (isWriteTool || isSafeReadOnlyTool) {
|
||||
const pathKey =
|
||||
'notebook_path' in input
|
||||
? 'notebook_path'
|
||||
: 'path' in input
|
||||
? 'path'
|
||||
: 'file_path'
|
||||
const filePath = input[pathKey] as string | undefined
|
||||
if (filePath) {
|
||||
const rel = relative(cwd, filePath)
|
||||
if (isAbsolute(rel) || rel.startsWith('..')) {
|
||||
if (isWriteTool) {
|
||||
logForDebugging(
|
||||
`[Speculation] Denied ${tool.name}: path outside cwd: ${filePath}`,
|
||||
)
|
||||
return denySpeculation(
|
||||
'Write outside cwd not allowed during speculation',
|
||||
'speculation_write_outside_root',
|
||||
)
|
||||
}
|
||||
return {
|
||||
behavior: 'allow' as const,
|
||||
updatedInput: input,
|
||||
decisionReason: {
|
||||
type: 'other' as const,
|
||||
reason: 'speculation_read_outside_root',
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
if (isWriteTool) {
|
||||
// Copy-on-write: copy original to overlay if not yet there
|
||||
if (!writtenPathsRef.current.has(rel)) {
|
||||
const overlayFile = join(overlayPath, rel)
|
||||
await mkdir(dirname(overlayFile), { recursive: true })
|
||||
try {
|
||||
await copyFile(join(cwd, rel), overlayFile)
|
||||
} catch {
|
||||
// Original may not exist (new file creation) - that's fine
|
||||
}
|
||||
writtenPathsRef.current.add(rel)
|
||||
}
|
||||
input = { ...input, [pathKey]: join(overlayPath, rel) }
|
||||
} else {
|
||||
// Read: redirect to overlay if file was previously written
|
||||
if (writtenPathsRef.current.has(rel)) {
|
||||
input = { ...input, [pathKey]: join(overlayPath, rel) }
|
||||
}
|
||||
// Otherwise read from main (no rewrite)
|
||||
}
|
||||
|
||||
logForDebugging(
|
||||
`[Speculation] ${isWriteTool ? 'Write' : 'Read'} ${filePath} -> ${input[pathKey]}`,
|
||||
)
|
||||
|
||||
return {
|
||||
behavior: 'allow' as const,
|
||||
updatedInput: input,
|
||||
decisionReason: {
|
||||
type: 'other' as const,
|
||||
reason: 'speculation_file_access',
|
||||
},
|
||||
}
|
||||
}
|
||||
// Read tools without explicit path (e.g. Glob/Grep defaulting to CWD) are safe
|
||||
if (isSafeReadOnlyTool) {
|
||||
return {
|
||||
behavior: 'allow' as const,
|
||||
updatedInput: input,
|
||||
decisionReason: {
|
||||
type: 'other' as const,
|
||||
reason: 'speculation_read_default_cwd',
|
||||
},
|
||||
}
|
||||
}
|
||||
// Write tools with undefined path → fall through to default deny
|
||||
}
|
||||
|
||||
// Stop at non-read-only bash commands
|
||||
if (tool.name === 'Bash') {
|
||||
const command =
|
||||
'command' in input && typeof input.command === 'string'
|
||||
? input.command
|
||||
: ''
|
||||
if (
|
||||
!command ||
|
||||
checkReadOnlyConstraints({ command }, commandHasAnyCd(command))
|
||||
.behavior !== 'allow'
|
||||
) {
|
||||
logForDebugging(
|
||||
`[Speculation] Stopping at bash: ${command.slice(0, 50) || 'missing command'}`,
|
||||
)
|
||||
updateActiveSpeculationState(setAppState, () => ({
|
||||
boundary: { type: 'bash', command, completedAt: Date.now() },
|
||||
}))
|
||||
abortController.abort()
|
||||
return denySpeculation(
|
||||
'Speculation paused: bash boundary',
|
||||
'speculation_bash_boundary',
|
||||
)
|
||||
}
|
||||
// Read-only bash command — allow during speculation
|
||||
return {
|
||||
behavior: 'allow' as const,
|
||||
updatedInput: input,
|
||||
decisionReason: {
|
||||
type: 'other' as const,
|
||||
reason: 'speculation_readonly_bash',
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Deny all other tools by default
|
||||
logForDebugging(`[Speculation] Stopping at denied tool: ${tool.name}`)
|
||||
const detail = String(
|
||||
('url' in input && input.url) ||
|
||||
('file_path' in input && input.file_path) ||
|
||||
('path' in input && input.path) ||
|
||||
('command' in input && input.command) ||
|
||||
'',
|
||||
).slice(0, 200)
|
||||
updateActiveSpeculationState(setAppState, () => ({
|
||||
boundary: {
|
||||
type: 'denied_tool',
|
||||
toolName: tool.name,
|
||||
detail,
|
||||
completedAt: Date.now(),
|
||||
},
|
||||
}))
|
||||
abortController.abort()
|
||||
return denySpeculation(
|
||||
`Tool ${tool.name} not allowed during speculation`,
|
||||
'speculation_unknown_tool',
|
||||
)
|
||||
},
|
||||
querySource: 'speculation',
|
||||
forkLabel: 'speculation',
|
||||
maxTurns: MAX_SPECULATION_TURNS,
|
||||
overrides: { abortController, requireCanUseTool: true },
|
||||
onMessage: msg => {
|
||||
if (msg.type === 'assistant' || msg.type === 'user') {
|
||||
messagesRef.current.push(msg)
|
||||
if (messagesRef.current.length >= MAX_SPECULATION_MESSAGES) {
|
||||
abortController.abort()
|
||||
}
|
||||
if (isUserMessageWithArrayContent(msg)) {
|
||||
const newTools = count(
|
||||
msg.message.content as { type: string; is_error?: boolean }[],
|
||||
b => b.type === 'tool_result' && !b.is_error,
|
||||
)
|
||||
if (newTools > 0) {
|
||||
updateActiveSpeculationState(setAppState, prev => ({
|
||||
toolUseCount: prev.toolUseCount + newTools,
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
})
|
||||
|
||||
if (abortController.signal.aborted) return
|
||||
|
||||
updateActiveSpeculationState(setAppState, () => ({
|
||||
boundary: {
|
||||
type: 'complete' as const,
|
||||
completedAt: Date.now(),
|
||||
outputTokens: result.totalUsage.output_tokens,
|
||||
},
|
||||
}))
|
||||
|
||||
logForDebugging(
|
||||
`[Speculation] Complete: ${countToolsInMessages(messagesRef.current)} tools`,
|
||||
)
|
||||
|
||||
// Pipeline: generate the next suggestion while we wait for the user to accept
|
||||
void generatePipelinedSuggestion(
|
||||
contextRef.current,
|
||||
suggestionText,
|
||||
messagesRef.current,
|
||||
setAppState,
|
||||
abortController,
|
||||
)
|
||||
} catch (error) {
|
||||
abortController.abort()
|
||||
|
||||
if (error instanceof Error && error.name === 'AbortError') {
|
||||
safeRemoveOverlay(overlayPath)
|
||||
resetSpeculationState(setAppState)
|
||||
return
|
||||
}
|
||||
|
||||
safeRemoveOverlay(overlayPath)
|
||||
|
||||
// eslint-disable-next-line no-restricted-syntax -- custom fallback message, not toError(e)
|
||||
logError(error instanceof Error ? error : new Error('Speculation failed'))
|
||||
|
||||
logSpeculation(
|
||||
id,
|
||||
'error',
|
||||
startTime,
|
||||
suggestionText.length,
|
||||
messagesRef.current,
|
||||
null,
|
||||
{
|
||||
error_type: error instanceof Error ? error.name : 'Unknown',
|
||||
error_message: errorMessage(error).slice(
|
||||
0,
|
||||
200,
|
||||
) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
error_phase:
|
||||
'start' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
is_pipelined: isPipelined,
|
||||
},
|
||||
)
|
||||
|
||||
resetSpeculationState(setAppState)
|
||||
}
|
||||
}
|
||||
|
||||
export async function acceptSpeculation(
|
||||
state: SpeculationState,
|
||||
setAppState: (f: (prev: AppState) => AppState) => void,
|
||||
cleanMessageCount: number,
|
||||
): Promise<SpeculationResult | null> {
|
||||
if (state.status !== 'active') return null
|
||||
|
||||
const {
|
||||
id,
|
||||
messagesRef,
|
||||
writtenPathsRef,
|
||||
abort,
|
||||
startTime,
|
||||
suggestionLength,
|
||||
isPipelined,
|
||||
} = state
|
||||
const messages = messagesRef.current
|
||||
const overlayPath = getOverlayPath(id)
|
||||
const acceptedAt = Date.now()
|
||||
|
||||
abort()
|
||||
|
||||
if (cleanMessageCount > 0) {
|
||||
await copyOverlayToMain(overlayPath, writtenPathsRef.current, getCwdState())
|
||||
}
|
||||
safeRemoveOverlay(overlayPath)
|
||||
|
||||
// Use snapshot boundary as default (available since state.status === 'active' was checked above)
|
||||
let boundary: CompletionBoundary | null = state.boundary
|
||||
let timeSavedMs =
|
||||
Math.min(acceptedAt, boundary?.completedAt ?? Infinity) - startTime
|
||||
|
||||
setAppState(prev => {
|
||||
// Refine with latest React state if speculation is still active
|
||||
if (prev.speculation.status === 'active' && prev.speculation.boundary) {
|
||||
boundary = prev.speculation.boundary
|
||||
const endTime = Math.min(acceptedAt, boundary.completedAt ?? Infinity)
|
||||
timeSavedMs = endTime - startTime
|
||||
}
|
||||
return {
|
||||
...prev,
|
||||
speculation: IDLE_SPECULATION_STATE,
|
||||
speculationSessionTimeSavedMs:
|
||||
prev.speculationSessionTimeSavedMs + timeSavedMs,
|
||||
}
|
||||
})
|
||||
|
||||
logForDebugging(
|
||||
boundary === null
|
||||
? `[Speculation] Accept ${id}: still running, using ${messages.length} messages`
|
||||
: `[Speculation] Accept ${id}: already complete`,
|
||||
)
|
||||
|
||||
logSpeculation(
|
||||
id,
|
||||
'accepted',
|
||||
startTime,
|
||||
suggestionLength,
|
||||
messages,
|
||||
boundary,
|
||||
{
|
||||
message_count: messages.length,
|
||||
time_saved_ms: timeSavedMs,
|
||||
is_pipelined: isPipelined,
|
||||
},
|
||||
)
|
||||
|
||||
if (timeSavedMs > 0) {
|
||||
const entry: SpeculationAcceptMessage = {
|
||||
type: 'speculation-accept',
|
||||
timestamp: new Date().toISOString(),
|
||||
timeSavedMs,
|
||||
}
|
||||
void appendFile(getTranscriptPath(), jsonStringify(entry) + '\n', {
|
||||
mode: 0o600,
|
||||
}).catch(() => {
|
||||
logForDebugging(
|
||||
'[Speculation] Failed to write speculation-accept to transcript',
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
return { messages, boundary, timeSavedMs }
|
||||
}
|
||||
|
||||
export function abortSpeculation(setAppState: SetAppState): void {
|
||||
setAppState(prev => {
|
||||
if (prev.speculation.status !== 'active') return prev
|
||||
|
||||
const {
|
||||
id,
|
||||
abort,
|
||||
startTime,
|
||||
boundary,
|
||||
suggestionLength,
|
||||
messagesRef,
|
||||
isPipelined,
|
||||
} = prev.speculation
|
||||
|
||||
logForDebugging(`[Speculation] Aborting ${id}`)
|
||||
|
||||
logSpeculation(
|
||||
id,
|
||||
'aborted',
|
||||
startTime,
|
||||
suggestionLength,
|
||||
messagesRef.current,
|
||||
boundary,
|
||||
{ abort_reason: 'user_typed', is_pipelined: isPipelined },
|
||||
)
|
||||
|
||||
abort()
|
||||
safeRemoveOverlay(getOverlayPath(id))
|
||||
|
||||
return { ...prev, speculation: IDLE_SPECULATION_STATE }
|
||||
})
|
||||
}
|
||||
|
||||
export async function handleSpeculationAccept(
|
||||
speculationState: ActiveSpeculationState,
|
||||
speculationSessionTimeSavedMs: number,
|
||||
setAppState: SetAppState,
|
||||
input: string,
|
||||
deps: {
|
||||
setMessages: (f: (prev: Message[]) => Message[]) => void
|
||||
readFileState: { current: FileStateCache }
|
||||
cwd: string
|
||||
},
|
||||
): Promise<{ queryRequired: boolean }> {
|
||||
try {
|
||||
const { setMessages, readFileState, cwd } = deps
|
||||
|
||||
// Clear prompt suggestion state. logOutcomeAtSubmission logged the accept
|
||||
// but was called with skipReset to avoid aborting speculation before we use it.
|
||||
setAppState(prev => {
|
||||
if (
|
||||
prev.promptSuggestion.text === null &&
|
||||
prev.promptSuggestion.promptId === null
|
||||
) {
|
||||
return prev
|
||||
}
|
||||
return {
|
||||
...prev,
|
||||
promptSuggestion: {
|
||||
text: null,
|
||||
promptId: null,
|
||||
shownAt: 0,
|
||||
acceptedAt: 0,
|
||||
generationRequestId: null,
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
// Capture speculation messages before any state updates - must be stable reference
|
||||
const speculationMessages = speculationState.messagesRef.current
|
||||
let cleanMessages = prepareMessagesForInjection(speculationMessages)
|
||||
|
||||
// Inject user message first for instant visual feedback before any async work
|
||||
const userMessage = createUserMessage({ content: input })
|
||||
setMessages(prev => [...prev, userMessage])
|
||||
|
||||
const result = await acceptSpeculation(
|
||||
speculationState,
|
||||
setAppState,
|
||||
cleanMessages.length,
|
||||
)
|
||||
|
||||
const isComplete = result?.boundary?.type === 'complete'
|
||||
|
||||
// When speculation didn't complete, the follow-up query needs the
|
||||
// conversation to end with a user message. Drop trailing assistant
|
||||
// messages — models that don't support prefill
|
||||
// reject conversations ending with an assistant turn. The model will
|
||||
// regenerate this content in the follow-up query.
|
||||
if (!isComplete) {
|
||||
const lastNonAssistant = cleanMessages.findLastIndex(
|
||||
m => m.type !== 'assistant',
|
||||
)
|
||||
cleanMessages = cleanMessages.slice(0, lastNonAssistant + 1)
|
||||
}
|
||||
|
||||
const timeSavedMs = result?.timeSavedMs ?? 0
|
||||
const newSessionTotal = speculationSessionTimeSavedMs + timeSavedMs
|
||||
const feedbackMessage = createSpeculationFeedbackMessage(
|
||||
cleanMessages,
|
||||
result?.boundary ?? null,
|
||||
timeSavedMs,
|
||||
newSessionTotal,
|
||||
)
|
||||
|
||||
// Inject speculated messages
|
||||
setMessages(prev => [...prev, ...cleanMessages])
|
||||
|
||||
const extracted = extractReadFilesFromMessages(
|
||||
cleanMessages,
|
||||
cwd,
|
||||
READ_FILE_STATE_CACHE_SIZE,
|
||||
)
|
||||
readFileState.current = mergeFileStateCaches(
|
||||
readFileState.current,
|
||||
extracted,
|
||||
)
|
||||
|
||||
if (feedbackMessage) {
|
||||
setMessages(prev => [...prev, feedbackMessage])
|
||||
}
|
||||
|
||||
logForDebugging(
|
||||
`[Speculation] ${result?.boundary?.type ?? 'incomplete'}, injected ${cleanMessages.length} messages`,
|
||||
)
|
||||
|
||||
// Promote pipelined suggestion if speculation completed fully
|
||||
if (isComplete && speculationState.pipelinedSuggestion) {
|
||||
const { text, promptId, generationRequestId } =
|
||||
speculationState.pipelinedSuggestion
|
||||
logForDebugging(
|
||||
`[Speculation] Promoting pipelined suggestion: "${text.slice(0, 50)}..."`,
|
||||
)
|
||||
setAppState(prev => ({
|
||||
...prev,
|
||||
promptSuggestion: {
|
||||
text,
|
||||
promptId,
|
||||
shownAt: Date.now(),
|
||||
acceptedAt: 0,
|
||||
generationRequestId,
|
||||
},
|
||||
}))
|
||||
|
||||
// Start speculation on the pipelined suggestion
|
||||
const augmentedContext: REPLHookContext = {
|
||||
...speculationState.contextRef.current,
|
||||
messages: [
|
||||
...speculationState.contextRef.current.messages,
|
||||
createUserMessage({ content: input }),
|
||||
...cleanMessages,
|
||||
],
|
||||
}
|
||||
void startSpeculation(text, augmentedContext, setAppState, true)
|
||||
}
|
||||
|
||||
return { queryRequired: !isComplete }
|
||||
} catch (error) {
|
||||
// Fail open: log error and fall back to normal query flow
|
||||
/* eslint-disable no-restricted-syntax -- custom fallback message, not toError(e) */
|
||||
logError(
|
||||
error instanceof Error
|
||||
? error
|
||||
: new Error('handleSpeculationAccept failed'),
|
||||
)
|
||||
/* eslint-enable no-restricted-syntax */
|
||||
logSpeculation(
|
||||
speculationState.id,
|
||||
'error',
|
||||
speculationState.startTime,
|
||||
speculationState.suggestionLength,
|
||||
speculationState.messagesRef.current,
|
||||
speculationState.boundary,
|
||||
{
|
||||
error_type: error instanceof Error ? error.name : 'Unknown',
|
||||
error_message: errorMessage(error).slice(
|
||||
0,
|
||||
200,
|
||||
) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
error_phase:
|
||||
'accept' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
is_pipelined: speculationState.isPipelined,
|
||||
},
|
||||
)
|
||||
safeRemoveOverlay(getOverlayPath(speculationState.id))
|
||||
resetSpeculationState(setAppState)
|
||||
// Query required so user's message is processed normally (without speculated work)
|
||||
return { queryRequired: true }
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user