diff --git a/.agents/types/tools.ts b/.agents/types/tools.ts index 754e54d78..15d036390 100644 --- a/.agents/types/tools.ts +++ b/.agents/types/tools.ts @@ -181,10 +181,10 @@ export interface ProposeStrReplaceParams { /** Array of replacements to make. */ replacements: { /** The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation. */ - old: string - /** The string to replace the corresponding old string with. Can be empty to delete. */ - new: string - /** Whether to allow multiple replacements of old string. */ + oldString: string + /** The string to replace the corresponding oldString with. Can be empty to delete. */ + newString: string + /** Whether to allow multiple replacements of oldString. */ allowMultiple?: boolean }[] } @@ -305,10 +305,10 @@ export interface StrReplaceParams { /** Array of replacements to make. */ replacements: { /** The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation. */ - old: string - /** The string to replace the corresponding old string with. Can be empty to delete. */ - new: string - /** Whether to allow multiple replacements of old string. */ + oldString: string + /** The string to replace the corresponding oldString with. Can be empty to delete. */ + newString: string + /** Whether to allow multiple replacements of oldString. */ allowMultiple?: boolean }[] } diff --git a/agents-graveyard/editor/reviewer-editor.ts b/agents-graveyard/editor/reviewer-editor.ts index c6cfe42b6..f76d8d559 100644 --- a/agents-graveyard/editor/reviewer-editor.ts +++ b/agents-graveyard/editor/reviewer-editor.ts @@ -36,12 +36,12 @@ Write out what changes you would make using the tool call format below. Use this "path": "path/to/file", "replacements": [ { - "old": "exact old code", - "new": "exact new code" + "oldString": "exact old code", + "newString": "exact new code" }, { - "old": "exact old code 2", - "new": "exact new code 2" + "oldString": "exact old code 2", + "newString": "exact new code 2" }, ] } diff --git a/agents/editor/best-of-n/editor-implementor.ts b/agents/editor/best-of-n/editor-implementor.ts index fe9fe13eb..2afc66d68 100644 --- a/agents/editor/best-of-n/editor-implementor.ts +++ b/agents/editor/best-of-n/editor-implementor.ts @@ -51,12 +51,12 @@ You can make multiple tool calls across multiple steps to complete the implement "path": "path/to/file", "replacements": [ { - "old": "exact old code", - "new": "exact new code" + "oldString": "exact old code", + "newString": "exact new code" }, { - "old": "exact old code 2", - "new": "exact new code 2" + "oldString": "exact old code 2", + "newString": "exact new code 2" }, ] } @@ -72,9 +72,10 @@ OR for new files or major rewrites: "content": "Complete file content" } -${isGpt5 || isGemini - ? `` - : ` +${ + isGpt5 || isGemini + ? `` + : ` IMPORTANT: Before you start writing your implementation, you should use tags to think about the best way to implement the changes. You should think really really hard to make sure you implement the changes in the best way possible. Take as much time as you to think through all the cases to produce the best changes. You can also use tags interspersed between tool calls to think about the best way to implement the changes. @@ -102,7 +103,7 @@ You can also use tags interspersed between tool calls to think about the ` - } +} After the edit tool calls, you can optionally mention any follow-up steps to take, like deleting a file, or a specific way to validate the changes. There's no need to use the set_output tool as your entire response will be included in the output. diff --git a/agents/editor/editor.ts b/agents/editor/editor.ts index 443724f67..a0cac064c 100644 --- a/agents/editor/editor.ts +++ b/agents/editor/editor.ts @@ -61,12 +61,12 @@ Write out what changes you would make using the tool call format below. Use this "path": "path/to/file", "replacements": [ { - "old": "exact old code", - "new": "exact new code" + "oldString": "exact old code", + "newString": "exact new code" }, { - "old": "exact old code 2", - "new": "exact new code 2" + "oldString": "exact old code 2", + "newString": "exact new code 2" }, ] } diff --git a/agents/types/tools.ts b/agents/types/tools.ts index 9cfe1cdf2..cb3882fc0 100644 --- a/agents/types/tools.ts +++ b/agents/types/tools.ts @@ -226,10 +226,10 @@ export interface ProposeStrReplaceParams { /** Array of replacements to make. */ replacements: { /** The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation. */ - old: string - /** The string to replace the corresponding old string with. Can be empty to delete. */ - new: string - /** Whether to allow multiple replacements of old string. */ + oldString: string + /** The string to replace the corresponding oldString with. Can be empty to delete. */ + newString: string + /** Whether to allow multiple replacements of oldString. */ allowMultiple?: boolean }[] } @@ -358,10 +358,10 @@ export interface StrReplaceParams { /** Array of replacements to make. */ replacements: { /** The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation. */ - old: string - /** The string to replace the corresponding old string with. Can be empty to delete. */ - new: string - /** Whether to allow multiple replacements of old string. */ + oldString: string + /** The string to replace the corresponding oldString with. Can be empty to delete. */ + newString: string + /** Whether to allow multiple replacements of oldString. */ allowMultiple?: boolean }[] } diff --git a/cli/src/components/freebuff-model-selector.tsx b/cli/src/components/freebuff-model-selector.tsx index 24f87350e..2552a1107 100644 --- a/cli/src/components/freebuff-model-selector.tsx +++ b/cli/src/components/freebuff-model-selector.tsx @@ -19,7 +19,10 @@ import { useFreebuffModelStore } from '../state/freebuff-model-store' import { useFreebuffSessionStore } from '../state/freebuff-session-store' import { useTerminalDimensions } from '../hooks/use-terminal-dimensions' import { useTheme } from '../hooks/use-theme' -import { nextFreebuffModelId } from '../utils/freebuff-model-navigation' +import { + freebuffModelNavigationDirectionForKey, + nextFreebuffModelId, +} from '../utils/freebuff-model-navigation' import type { FreebuffModelOption } from '@codebuff/common/constants/freebuff-models' import type { KeyEvent } from '@opentui/core' @@ -32,6 +35,9 @@ const FREEBUFF_MODEL_SELECTOR_MODELS: readonly FreebuffModelOption[] = [ ...FREEBUFF_MODELS.filter((model) => model.id === DEFAULT_FREEBUFF_MODEL_ID), ...FREEBUFF_MODELS.filter((model) => model.id !== DEFAULT_FREEBUFF_MODEL_ID), ] +const FREEBUFF_MODEL_SELECTOR_MODEL_IDS = FREEBUFF_MODEL_SELECTOR_MODELS.map( + (model) => model.id, +) function formatSessionUnits(units: number): string { return Number.isInteger(units) ? String(units) : units.toFixed(1) @@ -213,27 +219,26 @@ export const FreebuffModelSelector: React.FC = () => { (key: KeyEvent) => { if (pending) return const name = key.name ?? '' - const isForward = - name === 'right' || name === 'down' || (name === 'tab' && !key.shift) - const isBackward = - name === 'left' || name === 'up' || (name === 'tab' && key.shift) + const direction = freebuffModelNavigationDirectionForKey(key) const isCommit = name === 'return' || name === 'enter' || name === 'space' - if (!isForward && !isBackward && !isCommit) return if (isCommit) { if (isJoinable(focusedId) && focusedId !== committedModelId) { key.preventDefault?.() + key.stopPropagation?.() pick(focusedId) } return } + if (!direction) return const targetId = nextFreebuffModelId({ - modelIds: FREEBUFF_MODEL_SELECTOR_MODELS.map((model) => model.id), + modelIds: FREEBUFF_MODEL_SELECTOR_MODEL_IDS, focusedId, - direction: isForward ? 'forward' : 'backward', + direction, }) if (targetId) { key.preventDefault?.() + key.stopPropagation?.() setFocusedId(targetId) } }, diff --git a/cli/src/components/tools/str-replace.tsx b/cli/src/components/tools/str-replace.tsx index 881152472..ab1cc3823 100644 --- a/cli/src/components/tools/str-replace.tsx +++ b/cli/src/components/tools/str-replace.tsx @@ -3,43 +3,15 @@ import { TextAttributes } from '@opentui/core' import { DiffViewer } from './diff-viewer' import { defineToolComponent } from './types' import { useTheme } from '../../hooks/use-theme' +import { + extractDiff, + extractFilePath, + isCreateFile, + shouldShowEditDiff, +} from '../../utils/implementor-helpers' import type { ToolRenderConfig } from './types' -function extractValueForKey(output: string, key: string): string | null { - if (!output) return null - const lines = output.split('\n') - for (let i = 0; i < lines.length; i++) { - const line = lines[i] - const match = line.match(/^\s*([A-Za-z0-9_]+):\s*(.*)$/) - if (match && match[1] === key) { - const rest = match[2] - if (rest.trim().startsWith('|')) { - const baseIndent = lines[i + 1]?.match(/^\s*/)?.[0].length ?? 0 - const acc: string[] = [] - for (let j = i + 1; j < lines.length; j++) { - const l = lines[j] - const indent = l.match(/^\s*/)?.[0].length ?? 0 - if (l.trim().length === 0) { - acc.push('') - continue - } - if (indent < baseIndent) break - acc.push(l.slice(baseIndent)) - } - return acc.join('\n') - } else { - let val = rest.trim() - if (val.startsWith('"') && val.endsWith('"')) { - val = val.slice(1, -1) - } - return val - } - } - } - return null -} - interface EditHeaderProps { name: string filePath: string | null @@ -73,7 +45,7 @@ const EditBody = ({ name, filePath, diffText, isCreate }: EditBodyProps) => { return ( - {!isCreate && ( + {!isCreate && diffText.length > 0 && ( @@ -86,25 +58,17 @@ export const StrReplaceComponent = defineToolComponent({ toolName: 'str_replace', render(toolBlock): ToolRenderConfig { - const outputStr = - typeof toolBlock.output === 'string' ? toolBlock.output : '' - const diff = - extractValueForKey(outputStr, 'unifiedDiff') || - extractValueForKey(outputStr, 'patch') - const filePath = - extractValueForKey(outputStr, 'file') || - (typeof (toolBlock.input as any)?.path === 'string' - ? (toolBlock.input as any).path - : null) - const message = extractValueForKey(outputStr, 'message') - const isCreate = message === 'Created new file' + const diff = extractDiff(toolBlock) + const filePath = extractFilePath(toolBlock) + const isCreate = isCreateFile(toolBlock) + const showDiff = shouldShowEditDiff(toolBlock) return { content: ( ), diff --git a/cli/src/components/waiting-room-screen.tsx b/cli/src/components/waiting-room-screen.tsx index 9cdc385c9..8734bcaf1 100644 --- a/cli/src/components/waiting-room-screen.tsx +++ b/cli/src/components/waiting-room-screen.tsx @@ -93,8 +93,7 @@ const formatPrivacySignalList = ( const TakeoverPrompt: React.FC = () => { const theme = useTheme() const [pending, setPending] = useState(false) - const [takeoverHover, setTakeoverHover] = useState(false) - const [exitHover, setExitHover] = useState(false) + const [focusedIndex, setFocusedIndex] = useState(0) // 0 = Take over, 1 = Exit const handleTakeover = useCallback(() => { if (pending) return @@ -108,41 +107,79 @@ const TakeoverPrompt: React.FC = () => { const name = key.name ?? '' const isConfirm = name === 'return' || name === 'enter' const isExit = name === 'escape' || name === 'esc' - if (!isConfirm && !isExit) return - key.preventDefault?.() - if (isConfirm) { - handleTakeover() - } else { + const isTab = name === 'tab' + const isShiftTab = key.shift === true && isTab + const isRight = name === 'right' + const isLeft = name === 'left' + + if (isExit) { + key.preventDefault?.() exitFreebuffCleanly() + return + } + + if (isConfirm) { + key.preventDefault?.() + if (focusedIndex === 0) { + handleTakeover() + } else { + exitFreebuffCleanly() + } + return + } + + if (isRight || isTab) { + key.preventDefault?.() + setFocusedIndex((prev) => (prev + 1) % 2) + return + } + + if (isLeft || isShiftTab) { + key.preventDefault?.() + setFocusedIndex((prev) => (prev - 1 + 2) % 2) + return } }, - [handleTakeover], + [focusedIndex, handleTakeover], ), ) + const isTakeoverFocused = focusedIndex === 0 + const isExitFocused = focusedIndex === 1 + return ( - <> + Freebuff is already running - - Only one freebuff instance can run at a time. Take over the other - instance here, or exit and keep using the one already running. + + + Only one freebuff instance is allowed at a time. + - - Enter takes over · Esc exits - - + ) } @@ -258,7 +294,7 @@ export const WaitingRoomScreen: React.FC = ({ > ✕ @@ -419,7 +455,7 @@ export const WaitingRoomScreen: React.FC = ({ {/* Shared premium-session quota exhausted. Terminal for this run — the user can exit and come - back once the oldest session in the window rolls off. */} + back once the daily Pacific reset passes. */} {session?.status === 'rate_limited' && ( <> @@ -430,7 +466,7 @@ export const WaitingRoomScreen: React.FC = ({ {formatSessionUnits(session.recentCount)} of {session.limit} {' '} - premium sessions in the last 20 hours. Try again in{' '} + premium sessions today. Try again in{' '} {formatRetryAfter(session.retryAfterMs)} diff --git a/cli/src/utils/__tests__/freebuff-model-navigation.test.ts b/cli/src/utils/__tests__/freebuff-model-navigation.test.ts index 0df2a19a1..68157d71a 100644 --- a/cli/src/utils/__tests__/freebuff-model-navigation.test.ts +++ b/cli/src/utils/__tests__/freebuff-model-navigation.test.ts @@ -1,6 +1,9 @@ import { describe, expect, test } from 'bun:test' -import { nextFreebuffModelId } from '../freebuff-model-navigation' +import { + freebuffModelNavigationDirectionForKey, + nextFreebuffModelId, +} from '../freebuff-model-navigation' describe('nextFreebuffModelId', () => { test('moves to the next model when moving forward', () => { @@ -49,3 +52,51 @@ describe('nextFreebuffModelId', () => { ).toBeNull() }) }) + +describe('freebuffModelNavigationDirectionForKey', () => { + test('maps arrow keys to model navigation directions', () => { + expect(freebuffModelNavigationDirectionForKey({ name: 'down' })).toBe( + 'forward', + ) + expect(freebuffModelNavigationDirectionForKey({ name: 'right' })).toBe( + 'forward', + ) + expect(freebuffModelNavigationDirectionForKey({ name: 'up' })).toBe( + 'backward', + ) + expect(freebuffModelNavigationDirectionForKey({ name: 'left' })).toBe( + 'backward', + ) + }) + + test('maps tab and shift-tab to model navigation directions', () => { + expect(freebuffModelNavigationDirectionForKey({ name: 'tab' })).toBe( + 'forward', + ) + expect( + freebuffModelNavigationDirectionForKey({ name: 'tab', shift: true }), + ).toBe('backward') + }) + + test('maps terminal tab sequences to model navigation directions', () => { + expect(freebuffModelNavigationDirectionForKey({ sequence: '\t' })).toBe( + 'forward', + ) + expect( + freebuffModelNavigationDirectionForKey({ sequence: '\x1b[9u' }), + ).toBe('forward') + expect( + freebuffModelNavigationDirectionForKey({ sequence: '\x1b[Z' }), + ).toBe('backward') + expect( + freebuffModelNavigationDirectionForKey({ sequence: '\x1b[9;2u' }), + ).toBe('backward') + expect( + freebuffModelNavigationDirectionForKey({ sequence: '\x1b[27;2;9~' }), + ).toBe('backward') + }) + + test('ignores non-navigation keys', () => { + expect(freebuffModelNavigationDirectionForKey({ name: 'enter' })).toBeNull() + }) +}) diff --git a/cli/src/utils/__tests__/implementor-helpers.test.ts b/cli/src/utils/__tests__/implementor-helpers.test.ts index 83bcf2490..44793c408 100644 --- a/cli/src/utils/__tests__/implementor-helpers.test.ts +++ b/cli/src/utils/__tests__/implementor-helpers.test.ts @@ -17,9 +17,15 @@ import { groupConsecutiveToolBlocks, getMultiPromptProgress, getMultiPromptPreview, + shouldShowEditDiff, } from '../implementor-helpers' -import type { ToolContentBlock, ContentBlock, AgentContentBlock, TextContentBlock } from '../../types/chat' +import type { + ToolContentBlock, + ContentBlock, + AgentContentBlock, + TextContentBlock, +} from '../../types/chat' describe('extractValueForKey', () => { test('extracts simple key-value pairs', () => { @@ -104,9 +110,7 @@ describe('extractDiff', () => { toolCallId: 'test-1', toolName: 'str_replace', input: { - replacements: [ - { old: 'const x = 1', new: 'const x = 2' } - ] + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], }, } const diff = extractDiff(block) @@ -114,6 +118,82 @@ describe('extractDiff', () => { expect(diff).toContain('+ const x = 2') }) + test('constructs diff from successful str_replace input when output omits diff', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], + }, + output: 'message: String replace applied successfully.', + } + const diff = extractDiff(block) + expect(diff).toContain('- const x = 1') + expect(diff).toContain('+ const x = 2') + }) + + test('constructs diff from successful str_replace input with warning output', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], + }, + output: `message: | + Matched with indentation modification + + String replace applied successfully.`, + } + const diff = extractDiff(block) + expect(diff).toContain('- const x = 1') + expect(diff).toContain('+ const x = 2') + }) + + test('uses patch content from successful str_replace input when output omits diff', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { type: 'patch', content: '- const x = 1\n+ const x = 2' }, + output: 'message: String replace applied successfully.', + } + expect(extractDiff(block)).toBe('- const x = 1\n+ const x = 2') + }) + + test('returns null for failed str_replace output without a diff', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], + }, + output: 'No change to the file', + } + expect(extractDiff(block)).toBeNull() + }) + + test('returns null for failed str_replace output even when it includes patch input', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { type: 'patch', content: '- const x = 1\n+ const x = 2' }, + outputRaw: [ + { + type: 'json', + value: { + errorMessage: 'Failed to apply patch.', + patch: '- const x = 1\n+ const x = 2', + }, + }, + ], + } + expect(extractDiff(block)).toBeNull() + }) + test('constructs diff from write_file input', () => { const block: ToolContentBlock = { type: 'tool', @@ -125,15 +205,36 @@ describe('extractDiff', () => { expect(diff).toBe('+ line1\n+ line2') }) + test('constructs diff from successful write_file input when output omits diff', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: { content: 'line1\nline2' }, + output: 'message: Overwrote file successfully.', + } + const diff = extractDiff(block) + expect(diff).toBe('+ line1\n+ line2') + }) + + test('returns null for failed write_file output without a diff', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: { content: 'line1\nline2' }, + output: 'Failed to write to file', + } + expect(extractDiff(block)).toBeNull() + }) + test('constructs diff from propose_str_replace input', () => { const block: ToolContentBlock = { type: 'tool', toolCallId: 'test-1', toolName: 'propose_str_replace', input: { - replacements: [ - { old: 'const x = 1', new: 'const x = 2' } - ] + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], }, } const diff = extractDiff(block) @@ -178,8 +279,16 @@ describe('parseDiffStats', () => { }) test('handles empty diff', () => { - expect(parseDiffStats(undefined)).toEqual({ linesAdded: 0, linesRemoved: 0, hunks: 0 }) - expect(parseDiffStats('')).toEqual({ linesAdded: 0, linesRemoved: 0, hunks: 0 }) + expect(parseDiffStats(undefined)).toEqual({ + linesAdded: 0, + linesRemoved: 0, + hunks: 0, + }) + expect(parseDiffStats('')).toEqual({ + linesAdded: 0, + linesRemoved: 0, + hunks: 0, + }) }) test('ignores +++ and --- headers', () => { @@ -206,6 +315,17 @@ describe('getFileChangeType', () => { expect(getFileChangeType(block)).toBe('A') }) + test('returns A for successful file creation', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: {}, + output: 'message: Created file successfully.', + } + expect(getFileChangeType(block)).toBe('A') + }) + test('returns M for write_file modification', () => { const block: ToolContentBlock = { type: 'tool', @@ -249,6 +369,82 @@ describe('getFileChangeType', () => { }) }) +describe('shouldShowEditDiff', () => { + test('does not show pending str_replace diffs before the result arrives', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], + }, + } + + expect(shouldShowEditDiff(block)).toBe(false) + }) + + test('shows str_replace diffs after a successful result', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { + replacements: [{ oldString: 'const x = 1', newString: 'const x = 2' }], + }, + output: 'file: src/existing.ts\nmessage: String replace applied successfully.', + } + + expect(shouldShowEditDiff(block)).toBe(true) + }) + + test('does not show pending write_file diffs before the result arrives', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: { path: 'src/new.ts', content: 'const x = 1\n' }, + } + + expect(extractDiff(block)).toBe('+ const x = 1\n+ ') + expect(shouldShowEditDiff(block)).toBe(false) + }) + + test('shows write_file diffs after an overwrite result', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: { path: 'src/existing.ts', content: 'const x = 2\n' }, + output: 'file: src/existing.ts\nmessage: Overwrote file successfully.', + } + + expect(shouldShowEditDiff(block)).toBe(true) + }) + + test('does not show write_file diffs after a create result', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: { path: 'src/new.ts', content: 'const x = 1\n' }, + output: 'file: src/new.ts\nmessage: Created file successfully.', + } + + expect(shouldShowEditDiff(block)).toBe(false) + }) + + test('continues to show pending proposed write_file diffs', () => { + const block: ToolContentBlock = { + type: 'tool', + toolCallId: 'test-1', + toolName: 'propose_write_file', + input: { path: 'src/new.ts', content: 'const x = 1\n' }, + } + + expect(shouldShowEditDiff(block)).toBe(true) + }) +}) + describe('getFileStatsFromBlocks', () => { test('aggregates stats for same file', () => { const blocks: ContentBlock[] = [ @@ -264,7 +460,9 @@ describe('getFileStatsFromBlocks', () => { toolCallId: 'test-2', toolName: 'str_replace', input: { path: 'file.ts' }, - outputRaw: [{ type: 'json', value: { unifiedDiff: '+line3\n-removed' } }], + outputRaw: [ + { type: 'json', value: { unifiedDiff: '+line3\n-removed' } }, + ], }, ] const stats = getFileStatsFromBlocks(blocks) @@ -307,6 +505,25 @@ describe('getFileStatsFromBlocks', () => { const stats = getFileStatsFromBlocks(blocks) expect(stats).toHaveLength(0) }) + + test('ignores failed edit tools', () => { + const blocks: ContentBlock[] = [ + { + type: 'tool', + toolCallId: 'test-1', + toolName: 'str_replace', + input: { + path: 'file.ts', + replacements: [ + { oldString: 'const x = 1', newString: 'const x = 2' }, + ], + }, + output: 'No change to the file', + }, + ] + const stats = getFileStatsFromBlocks(blocks) + expect(stats).toHaveLength(0) + }) }) describe('buildActivityTimeline', () => { @@ -354,20 +571,53 @@ describe('buildActivityTimeline', () => { expect(timeline).toHaveLength(1) expect(timeline[0].content).toBe('Normal text') }) + + test('skips failed edit tools', () => { + const blocks: ContentBlock[] = [ + { + type: 'text', + content: 'Trying an edit', + } as TextContentBlock, + { + type: 'tool', + toolCallId: 'test-1', + toolName: 'write_file', + input: { path: 'file.ts', content: 'new content' }, + output: 'Failed to write to file', + }, + ] + const timeline = buildActivityTimeline(blocks) + expect(timeline).toHaveLength(1) + expect(timeline[0].type).toBe('commentary') + }) }) describe('isImplementorAgent', () => { test('identifies implementor agents', () => { - expect(isImplementorAgent({ agentType: 'editor-implementor', blocks: [] })).toBe(true) - expect(isImplementorAgent({ agentType: 'editor-implementor-opus', blocks: [] })).toBe(true) - expect(isImplementorAgent({ agentType: 'editor-implementor-gpt-5', blocks: [] })).toBe(true) - expect(isImplementorAgent({ agentType: 'editor-implementor2', blocks: [] })).toBe(true) + expect( + isImplementorAgent({ agentType: 'editor-implementor', blocks: [] }), + ).toBe(true) + expect( + isImplementorAgent({ agentType: 'editor-implementor-opus', blocks: [] }), + ).toBe(true) + expect( + isImplementorAgent({ agentType: 'editor-implementor-gpt-5', blocks: [] }), + ).toBe(true) + expect( + isImplementorAgent({ agentType: 'editor-implementor2', blocks: [] }), + ).toBe(true) }) test('rejects non-implementor agents', () => { - expect(isImplementorAgent({ agentType: 'file-picker', blocks: [] })).toBe(false) - expect(isImplementorAgent({ agentType: 'commander', blocks: [] })).toBe(false) - expect(isImplementorAgent({ agentType: 'best-of-n-selector', blocks: [] })).toBe(false) + expect(isImplementorAgent({ agentType: 'file-picker', blocks: [] })).toBe( + false, + ) + expect(isImplementorAgent({ agentType: 'commander', blocks: [] })).toBe( + false, + ) + expect( + isImplementorAgent({ agentType: 'best-of-n-selector', blocks: [] }), + ).toBe(false) }) }) @@ -376,20 +626,48 @@ describe('getImplementorDisplayName', () => { expect(getImplementorDisplayName('editor-implementor')).toBe('Sonnet') expect(getImplementorDisplayName('editor-implementor-opus')).toBe('Opus') expect(getImplementorDisplayName('editor-implementor-gpt-5')).toBe('GPT-5') - expect(getImplementorDisplayName('editor-implementor-gemini')).toBe('Gemini') + expect(getImplementorDisplayName('editor-implementor-gemini')).toBe( + 'Gemini', + ) }) test('adds index when provided', () => { expect(getImplementorDisplayName('editor-implementor', 0)).toBe('Sonnet #1') - expect(getImplementorDisplayName('editor-implementor-opus', 2)).toBe('Opus #3') + expect(getImplementorDisplayName('editor-implementor-opus', 2)).toBe( + 'Opus #3', + ) }) }) describe('getImplementorIndex', () => { test('returns index among same-type siblings', () => { - const agent1 = { type: 'agent', agentId: 'a1', agentName: 'Impl 1', agentType: 'editor-implementor', content: '', status: 'complete', blocks: [] } as AgentContentBlock - const agent2 = { type: 'agent', agentId: 'a2', agentName: 'Impl 2', agentType: 'editor-implementor', content: '', status: 'complete', blocks: [] } as AgentContentBlock - const agent3 = { type: 'agent', agentId: 'a3', agentName: 'Impl 3', agentType: 'editor-implementor-opus', content: '', status: 'complete', blocks: [] } as AgentContentBlock + const agent1 = { + type: 'agent', + agentId: 'a1', + agentName: 'Impl 1', + agentType: 'editor-implementor', + content: '', + status: 'complete', + blocks: [], + } as AgentContentBlock + const agent2 = { + type: 'agent', + agentId: 'a2', + agentName: 'Impl 2', + agentType: 'editor-implementor', + content: '', + status: 'complete', + blocks: [], + } as AgentContentBlock + const agent3 = { + type: 'agent', + agentId: 'a3', + agentName: 'Impl 3', + agentType: 'editor-implementor-opus', + content: '', + status: 'complete', + blocks: [], + } as AgentContentBlock const siblings: ContentBlock[] = [agent1, agent2, agent3] expect(getImplementorIndex(agent1, siblings)).toBe(0) @@ -398,7 +676,15 @@ describe('getImplementorIndex', () => { }) test('returns undefined for non-implementor', () => { - const filePicker = { type: 'agent', agentId: 'fp1', agentName: 'File Picker', agentType: 'file-picker', content: '', status: 'complete', blocks: [] } as AgentContentBlock + const filePicker = { + type: 'agent', + agentId: 'fp1', + agentName: 'File Picker', + agentType: 'file-picker', + content: '', + status: 'complete', + blocks: [], + } as AgentContentBlock const siblings: ContentBlock[] = [filePicker] expect(getImplementorIndex(filePicker, siblings)).toBeUndefined() @@ -406,10 +692,11 @@ describe('getImplementorIndex', () => { }) describe('groupConsecutiveBlocks', () => { - const createTextBlock = (content: string): TextContentBlock => ({ - type: 'text', - content, - } as TextContentBlock) + const createTextBlock = (content: string): TextContentBlock => + ({ + type: 'text', + content, + }) as TextContentBlock const createToolBlock = (toolName: string): ToolContentBlock => ({ type: 'tool', @@ -418,15 +705,19 @@ describe('groupConsecutiveBlocks', () => { input: {}, }) - const createAgentBlock = (agentType: string, agentId: string): AgentContentBlock => ({ - type: 'agent', - agentId, - agentName: agentType, - agentType, - content: '', - status: 'complete', - blocks: [], - } as AgentContentBlock) + const createAgentBlock = ( + agentType: string, + agentId: string, + ): AgentContentBlock => + ({ + type: 'agent', + agentId, + agentName: agentType, + agentType, + content: '', + status: 'complete', + blocks: [], + }) as AgentContentBlock test('groups consecutive matching blocks from start', () => { const blocks: ContentBlock[] = [ @@ -530,7 +821,8 @@ describe('groupConsecutiveBlocks', () => { createTextBlock('done'), ] const isEditTool = (b: ContentBlock): b is ToolContentBlock => - b.type === 'tool' && ['str_replace', 'write_file'].includes(b.toolName as string) + b.type === 'tool' && + ['str_replace', 'write_file'].includes(b.toolName as string) const result = groupConsecutiveBlocks(blocks, 0, isEditTool) expect(result.group).toHaveLength(2) @@ -541,30 +833,39 @@ describe('groupConsecutiveBlocks', () => { }) describe('groupConsecutiveImplementors', () => { - const createImplementorAgent = (id: string, agentType = 'editor-implementor'): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: 'Implementor', - agentType, - content: '', - status: 'complete', - blocks: [], - } as AgentContentBlock) - - const createNonImplementorAgent = (id: string, agentType: string): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: agentType, - agentType, - content: '', - status: 'complete', - blocks: [], - } as AgentContentBlock) - - const createTextBlock = (content: string): TextContentBlock => ({ - type: 'text', - content, - } as TextContentBlock) + const createImplementorAgent = ( + id: string, + agentType = 'editor-implementor', + ): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: 'Implementor', + agentType, + content: '', + status: 'complete', + blocks: [], + }) as AgentContentBlock + + const createNonImplementorAgent = ( + id: string, + agentType: string, + ): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: agentType, + agentType, + content: '', + status: 'complete', + blocks: [], + }) as AgentContentBlock + + const createTextBlock = (content: string): TextContentBlock => + ({ + type: 'text', + content, + }) as TextContentBlock test('groups consecutive implementor agents', () => { const blocks: ContentBlock[] = [ @@ -654,30 +955,36 @@ describe('groupConsecutiveImplementors', () => { }) describe('groupConsecutiveNonImplementorAgents', () => { - const createImplementorAgent = (id: string): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: 'Implementor', - agentType: 'editor-implementor', - content: '', - status: 'complete', - blocks: [], - } as AgentContentBlock) - - const createNonImplementorAgent = (id: string, agentType: string): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: agentType, - agentType, - content: '', - status: 'complete', - blocks: [], - } as AgentContentBlock) - - const createTextBlock = (content: string): TextContentBlock => ({ - type: 'text', - content, - } as TextContentBlock) + const createImplementorAgent = (id: string): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: 'Implementor', + agentType: 'editor-implementor', + content: '', + status: 'complete', + blocks: [], + }) as AgentContentBlock + + const createNonImplementorAgent = ( + id: string, + agentType: string, + ): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: agentType, + agentType, + content: '', + status: 'complete', + blocks: [], + }) as AgentContentBlock + + const createTextBlock = (content: string): TextContentBlock => + ({ + type: 'text', + content, + }) as TextContentBlock test('groups consecutive non-implementor agents', () => { const blocks: ContentBlock[] = [ @@ -776,25 +1083,32 @@ describe('groupConsecutiveNonImplementorAgents', () => { }) describe('getMultiPromptProgress', () => { - const createImplementorAgent = (id: string, status: 'running' | 'complete' | 'failed' | 'cancelled' = 'complete'): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: 'Implementor', - agentType: 'editor-implementor-opus', - content: '', - status, - blocks: [], - } as AgentContentBlock) - - const createSelectorAgent = (status: 'running' | 'complete' = 'running'): AgentContentBlock => ({ - type: 'agent', - agentId: 'selector-1', - agentName: 'Selector', - agentType: 'best-of-n-selector2', - content: '', - status, - blocks: [], - } as AgentContentBlock) + const createImplementorAgent = ( + id: string, + status: 'running' | 'complete' | 'failed' | 'cancelled' = 'complete', + ): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: 'Implementor', + agentType: 'editor-implementor-opus', + content: '', + status, + blocks: [], + }) as AgentContentBlock + + const createSelectorAgent = ( + status: 'running' | 'complete' = 'running', + ): AgentContentBlock => + ({ + type: 'agent', + agentId: 'selector-1', + agentName: 'Selector', + agentType: 'best-of-n-selector2', + content: '', + status, + blocks: [], + }) as AgentContentBlock test('returns null for empty blocks', () => { expect(getMultiPromptProgress([])).toBeNull() @@ -877,31 +1191,40 @@ describe('getMultiPromptProgress', () => { }) describe('getMultiPromptPreview', () => { - const createImplementorAgent = (id: string, status: 'running' | 'complete' | 'failed' | 'cancelled' = 'complete'): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: 'Implementor', - agentType: 'editor-implementor-opus', - content: '', - status, - blocks: [], - } as AgentContentBlock) - - const createSelectorAgent = (status: 'running' | 'complete' = 'running'): AgentContentBlock => ({ - type: 'agent', - agentId: 'selector-1', - agentName: 'Selector', - agentType: 'best-of-n-selector2', - content: '', - status, - blocks: [], - } as AgentContentBlock) + const createImplementorAgent = ( + id: string, + status: 'running' | 'complete' | 'failed' | 'cancelled' = 'complete', + ): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: 'Implementor', + agentType: 'editor-implementor-opus', + content: '', + status, + blocks: [], + }) as AgentContentBlock + + const createSelectorAgent = ( + status: 'running' | 'complete' = 'running', + ): AgentContentBlock => + ({ + type: 'agent', + agentId: 'selector-1', + agentName: 'Selector', + agentType: 'best-of-n-selector2', + content: '', + status, + blocks: [], + }) as AgentContentBlock const createSetOutputBlock = (reason?: string): ToolContentBlock => ({ type: 'tool', toolCallId: 'set-output-1', toolName: 'set_output', - input: reason ? { data: { chosenStrategy: 'strategy A', reason } } : { data: { chosenStrategy: 'strategy A' } }, + input: reason + ? { data: { chosenStrategy: 'strategy A', reason } } + : { data: { chosenStrategy: 'strategy A' } }, }) test('returns null for empty blocks', () => { @@ -934,7 +1257,9 @@ describe('getMultiPromptPreview', () => { createImplementorAgent('impl-3', 'complete'), createSelectorAgent('running'), ] - expect(getMultiPromptPreview(blocks)).toBe('3 proposals complete • Selecting best...') + expect(getMultiPromptPreview(blocks)).toBe( + '3 proposals complete • Selecting best...', + ) }) test('shows applying message when selector is complete but agent not done', () => { @@ -943,7 +1268,9 @@ describe('getMultiPromptPreview', () => { createImplementorAgent('impl-2', 'complete'), createSelectorAgent('complete'), ] - expect(getMultiPromptPreview(blocks, false)).toBe('Applying selected changes...') + expect(getMultiPromptPreview(blocks, false)).toBe( + 'Applying selected changes...', + ) }) test('shows evaluation count when agent is complete without reason', () => { @@ -962,7 +1289,9 @@ describe('getMultiPromptPreview', () => { createSetOutputBlock('best implementation with proper error handling'), ] const preview = getMultiPromptPreview(blocks, true) - expect(preview).toBe('2 proposals evaluated\nBest implementation with proper error handling') + expect(preview).toBe( + '2 proposals evaluated\nBest implementation with proper error handling', + ) }) test('capitalizes first letter of reason', () => { @@ -989,7 +1318,9 @@ describe('getMultiPromptPreview', () => { createImplementorAgent('impl-2', 'complete'), createImplementorAgent('impl-3', 'failed'), ] - expect(getMultiPromptPreview(blocks)).toBe('2/3 proposals complete (1 failed)') + expect(getMultiPromptPreview(blocks)).toBe( + '2/3 proposals complete (1 failed)', + ) }) test('treats failed implementors as finished for progress', () => { @@ -999,7 +1330,9 @@ describe('getMultiPromptPreview', () => { createImplementorAgent('impl-3', 'complete'), ] // All 3 are finished (1 complete + 2 failed/cancelled), so should show completion message - expect(getMultiPromptPreview(blocks)).toBe('1/3 proposals complete (2 failed)') + expect(getMultiPromptPreview(blocks)).toBe( + '1/3 proposals complete (2 failed)', + ) }) }) @@ -1011,20 +1344,22 @@ describe('groupConsecutiveToolBlocks', () => { input: {}, }) - const createTextBlock = (content: string): TextContentBlock => ({ - type: 'text', - content, - } as TextContentBlock) - - const createAgentBlock = (id: string): AgentContentBlock => ({ - type: 'agent', - agentId: id, - agentName: 'Test Agent', - agentType: 'file-picker', - content: '', - status: 'complete', - blocks: [], - } as AgentContentBlock) + const createTextBlock = (content: string): TextContentBlock => + ({ + type: 'text', + content, + }) as TextContentBlock + + const createAgentBlock = (id: string): AgentContentBlock => + ({ + type: 'agent', + agentId: id, + agentName: 'Test Agent', + agentType: 'file-picker', + content: '', + status: 'complete', + blocks: [], + }) as AgentContentBlock test('groups consecutive tool blocks', () => { const blocks: ContentBlock[] = [ diff --git a/cli/src/utils/freebuff-model-navigation.ts b/cli/src/utils/freebuff-model-navigation.ts index d1f748d8c..a866ae16a 100644 --- a/cli/src/utils/freebuff-model-navigation.ts +++ b/cli/src/utils/freebuff-model-navigation.ts @@ -1,7 +1,18 @@ +export type FreebuffModelNavigationDirection = 'forward' | 'backward' + +const FORWARD_KEY_NAMES = new Set(['right', 'down']) +const BACKWARD_KEY_NAMES = new Set(['left', 'up']) +const FORWARD_TAB_SEQUENCES = new Set(['\t', '\x1b[9u']) +const BACKWARD_TAB_SEQUENCES = new Set([ + '\x1b[Z', + '\x1b[9;2u', + '\x1b[27;2;9~', +]) + export function nextFreebuffModelId(params: { modelIds: readonly string[] focusedId: string - direction: 'forward' | 'backward' + direction: FreebuffModelNavigationDirection }): string | null { const { modelIds, focusedId, direction } = params if (modelIds.length === 0) return null @@ -12,3 +23,28 @@ export function nextFreebuffModelId(params: { const step = direction === 'forward' ? 1 : -1 return modelIds[(currentIdx + step + modelIds.length) % modelIds.length] } + +export function freebuffModelNavigationDirectionForKey(key: { + name?: string + shift?: boolean + sequence?: string + raw?: string +}): FreebuffModelNavigationDirection | null { + const name = (key.name ?? '').toLowerCase() + const sequence = key.sequence ?? key.raw ?? '' + + if (FORWARD_KEY_NAMES.has(name)) return 'forward' + if (BACKWARD_KEY_NAMES.has(name)) return 'backward' + + if ( + (name === 'tab' && Boolean(key.shift)) || + BACKWARD_TAB_SEQUENCES.has(sequence) + ) { + return 'backward' + } + if (name === 'tab' || FORWARD_TAB_SEQUENCES.has(sequence)) { + return 'forward' + } + + return null +} diff --git a/cli/src/utils/implementor-helpers.ts b/cli/src/utils/implementor-helpers.ts index ca757ba52..ccb92c5c1 100644 --- a/cli/src/utils/implementor-helpers.ts +++ b/cli/src/utils/implementor-helpers.ts @@ -25,6 +25,18 @@ const isProposedToolName = (toolName: ToolContentBlock['toolName']): boolean => const getBaseToolName = (toolName: ToolContentBlock['toolName']): string => isProposedToolName(toolName) ? toolName.slice('propose_'.length) : toolName +const SUCCESSFUL_EDIT_MESSAGES = [ + 'String replace applied successfully', + 'Created file successfully', + 'Created new file', + 'Overwrote file successfully', + 'Wrote file successfully', + 'Updated file', + 'Proposed new file', + 'Proposed changes', + 'Proposed string replacement', +] as const + const hasProposedTools = (blocks?: ContentBlock[]): boolean => { if (!blocks || blocks.length === 0) return false @@ -221,38 +233,61 @@ export function extractFilePath(toolBlock: ToolContentBlock): string | null { * For proposed tools (implementors): construct diff from input replacements. */ export function extractDiff(toolBlock: ToolContentBlock): string | null { + let hasSuccessfulOutput = false + // First try to get from outputRaw (for executed tool results) // outputRaw is typically an array like [{type: "json", value: {unifiedDiff: "..."}}] const outputRaw = toolBlock.outputRaw as unknown if (Array.isArray(outputRaw) && outputRaw[0]?.value) { const value = outputRaw[0].value as Record + if (hasErrorMessage(value)) return null + if (isSuccessfulEditMessage(value.message)) hasSuccessfulOutput = true if (value.unifiedDiff) return value.unifiedDiff as string if (value.patch) return value.patch as string } // Also check direct properties (in case format differs) if (typeof outputRaw === 'object' && outputRaw !== null) { const rawObj = outputRaw as Record + if (hasErrorMessage(rawObj)) return null + if (isSuccessfulEditMessage(rawObj.message)) hasSuccessfulOutput = true if (rawObj.unifiedDiff) return rawObj.unifiedDiff as string if (rawObj.patch) return rawObj.patch as string } // Try to get from output string (key: value format) const outputStr = typeof toolBlock.output === 'string' ? toolBlock.output : '' + const message = extractValueForKey(outputStr, 'message') const diffFromOutput = extractValueForKey(outputStr, 'unifiedDiff') || extractValueForKey(outputStr, 'patch') + if (hasFailedEditOutput({ outputStr, message, diffFromOutput })) { + return null + } + if (isSuccessfulEditMessage(message)) { + hasSuccessfulOutput = true + } + if (diffFromOutput) { return diffFromOutput } - // For proposed edits (no output yet): construct diff from input + // For proposed/pending edits, or confirmed successful executions, construct + // the preview from input when the result omits a diff. + const canUseInputFallback = + isProposedToolName(toolBlock.toolName) || + outputStr === '' || + hasSuccessfulOutput + if (!canUseInputFallback) { + return null + } + const input = toolBlock.input as Record const baseToolName = getBaseToolName(toolBlock.toolName) // Handle str_replace: construct diff from replacements if (baseToolName === 'str_replace' && Array.isArray(input?.replacements)) { - const replacements = input.replacements as { old: string; new: string }[] + const replacements = input.replacements as ReplacementInput[] if (replacements.length > 0) { return constructDiffFromReplacements(replacements) } @@ -271,22 +306,96 @@ export function extractDiff(toolBlock: ToolContentBlock): string | null { return null } +function hasErrorMessage(value: Record): boolean { + return Boolean(value.errorMessage || (value.value as any)?.errorMessage) +} + +function hasFailedEditOutput(params: { + outputStr: string + message: string | null + diffFromOutput: string | null +}): boolean { + const { outputStr, message, diffFromOutput } = params + const trimmedOutput = outputStr.trim() + if (!trimmedOutput) { + return false + } + if ( + extractValueForKey(outputStr, 'errorMessage') || + isErrorOutput(outputStr) + ) { + return true + } + if (diffFromOutput || isSuccessfulEditMessage(message)) { + return false + } + return !isSuccessfulEditMessage(trimmedOutput) +} + +function isFailedEditToolBlock(toolBlock: ToolContentBlock): boolean { + const outputRaw = toolBlock.outputRaw as unknown + if (Array.isArray(outputRaw) && outputRaw[0]?.value) { + const value = outputRaw[0].value as Record + if (hasErrorMessage(value)) return true + } + if (typeof outputRaw === 'object' && outputRaw !== null) { + const rawObj = outputRaw as Record + if (hasErrorMessage(rawObj)) return true + } + + const outputStr = typeof toolBlock.output === 'string' ? toolBlock.output : '' + const message = extractValueForKey(outputStr, 'message') + const diffFromOutput = + extractValueForKey(outputStr, 'unifiedDiff') || + extractValueForKey(outputStr, 'patch') + return hasFailedEditOutput({ outputStr, message, diffFromOutput }) +} + +function isSuccessfulEditMessage(message: unknown): boolean { + if (typeof message !== 'string') { + return false + } + + return message + .split('\n') + .some((line) => + SUCCESSFUL_EDIT_MESSAGES.some((successMessage) => + line.trim().startsWith(successMessage), + ), + ) +} + +function isErrorOutput(output: string): boolean { + const trimmedOutput = output.trim() + return trimmedOutput.startsWith('Error:') || trimmedOutput.startsWith('Failed ') +} + /** * Construct a simple diff view from str_replace replacements. */ +type ReplacementInput = { + oldString?: string + newString?: string + old?: string + new?: string +} + function constructDiffFromReplacements( - replacements: { old: string; new: string }[], + replacements: ReplacementInput[], ): string { const lines: string[] = [] for (const replacement of replacements) { + const oldString = replacement.oldString ?? replacement.old ?? '' + const newString = replacement.newString ?? replacement.new ?? '' + // Add old lines as removals - const oldLines = replacement.old.split('\n') + const oldLines = oldString.split('\n') for (const line of oldLines) { lines.push(`- ${line}`) } // Add new lines as additions - const newLines = replacement.new.split('\n') + const newLines = newString.split('\n') for (const line of newLines) { lines.push(`+ ${line}`) } @@ -315,11 +424,39 @@ export function isCreateFile(toolBlock: ToolContentBlock): boolean { const message = extractValueForKey(outputStr, 'message') return ( typeof message === 'string' && - (message.startsWith('Created new file') || + (message.startsWith('Created file successfully') || + message.startsWith('Created new file') || message.startsWith('Proposed new file')) ) } +function hasToolResultOutput(toolBlock: ToolContentBlock): boolean { + const outputStr = typeof toolBlock.output === 'string' ? toolBlock.output : '' + return outputStr.length > 0 || toolBlock.outputRaw !== undefined +} + +/** + * Decide whether the direct edit tool renderer should show a diff preview. + * + * Real edit tool calls render immediately with input only, then receive output + * once the edit completes. Wait for that result before showing diffs so create + * operations never briefly flash an input-derived full-file diff. + */ +export function shouldShowEditDiff(toolBlock: ToolContentBlock): boolean { + if (!extractDiff(toolBlock) || isCreateFile(toolBlock)) { + return false + } + + if ( + !isProposedToolName(toolBlock.toolName) && + !hasToolResultOutput(toolBlock) + ) { + return false + } + + return true +} + export interface TimelineItem { type: 'commentary' | 'edit' content: string // For commentary: the text. For edits: file path @@ -400,7 +537,9 @@ export function getFileChangeType(toolBlock: ToolContentBlock): FileChangeType { * Get aggregated file stats from all edit blocks. * Groups by file path and sums up the stats. */ -export function getFileStatsFromBlocks(blocks: ContentBlock[] | undefined): FileStats[] { +export function getFileStatsFromBlocks( + blocks: ContentBlock[] | undefined, +): FileStats[] { if (!blocks || blocks.length === 0) return [] const fileMap = new Map() @@ -408,8 +547,12 @@ export function getFileStatsFromBlocks(blocks: ContentBlock[] | undefined): File for (const block of blocks) { if ( block.type === 'tool' && - ALL_EDIT_TOOL_NAMES.includes(block.toolName as (typeof ALL_EDIT_TOOL_NAMES)[number]) + ALL_EDIT_TOOL_NAMES.includes( + block.toolName as (typeof ALL_EDIT_TOOL_NAMES)[number], + ) ) { + if (isFailedEditToolBlock(block)) continue + const filePath = extractFilePath(block) if (!filePath) continue @@ -456,8 +599,12 @@ export function buildActivityTimeline( } } else if ( block.type === 'tool' && - ALL_EDIT_TOOL_NAMES.includes(block.toolName as (typeof ALL_EDIT_TOOL_NAMES)[number]) + ALL_EDIT_TOOL_NAMES.includes( + block.toolName as (typeof ALL_EDIT_TOOL_NAMES)[number], + ) ) { + if (isFailedEditToolBlock(block)) continue + const filePath = extractFilePath(block) const diff = extractDiff(block) const isCreate = isCreateFile(block) @@ -519,8 +666,7 @@ export function getMultiPromptProgress( const selectorAgent = blocks.find( (block): block is AgentContentBlock => - block.type === 'agent' && - block.agentType.includes('best-of-n-selector'), + block.type === 'agent' && block.agentType.includes('best-of-n-selector'), ) const isSelecting = selectorAgent?.status === 'running' @@ -562,7 +708,9 @@ function hasSetOutputData(input: unknown): input is SetOutputInput { * Extract the selection reason from multi-prompt agent's set_output block. * set_output wraps data in a 'data' property, so we need to access input.data.reason */ -function extractSelectionReason(blocks: ContentBlock[] | undefined): string | null { +function extractSelectionReason( + blocks: ContentBlock[] | undefined, +): string | null { if (!blocks || blocks.length === 0) return null const setOutputBlock = blocks.find( @@ -604,7 +752,9 @@ export function getMultiPromptPreview( const formattedReason = reason.charAt(0).toUpperCase() + reason.slice(1) const lines = formattedReason.split('\n') const truncatedReason = - lines.length > 2 ? lines.slice(0, 2).join('\n').trimEnd() + '...' : formattedReason + lines.length > 2 + ? lines.slice(0, 2).join('\n').trimEnd() + '...' + : formattedReason return `${total} proposals evaluated\n${truncatedReason}` } return `${total} proposals evaluated` diff --git a/common/src/constants/freebuff-models.ts b/common/src/constants/freebuff-models.ts index fedd5154c..8bfaf7b76 100644 --- a/common/src/constants/freebuff-models.ts +++ b/common/src/constants/freebuff-models.ts @@ -1,3 +1,10 @@ +import { + addDaysToYmd, + getUtcForZonedTime, + getZonedParts, + type ZonedDateParts, +} from '../util/zoned-time' + /** * Models a freebuff user can pick between in the waiting-room model selector. * @@ -31,18 +38,14 @@ export const FREEBUFF_GLM_MODEL_ID = 'z-ai/glm-5.1' export const FREEBUFF_KIMI_MODEL_ID = 'moonshotai/kimi-k2.6' export const FREEBUFF_MINIMAX_MODEL_ID = 'minimax/minimax-m2.7' export const FREEBUFF_PREMIUM_SESSION_LIMIT = 5 -export const FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS = 20 +export const FREEBUFF_PREMIUM_SESSION_RESET_TIMEZONE = 'America/Los_Angeles' +export const FREEBUFF_PREMIUM_SESSION_PERIOD = 'pacific_day' +/** Deprecated wire compatibility field. Premium usage now resets at midnight + * Pacific time rather than using a rolling hourly window. */ +export const FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS = 24 const FREEBUFF_EASTERN_TIMEZONE = 'America/New_York' const FREEBUFF_PACIFIC_TIMEZONE = 'America/Los_Angeles' -interface ZonedDateParts { - year: number - month: number - day: number - hour: number - minute: number -} - interface LocalTimeFormatOptions { locale?: string timeZone?: string @@ -165,79 +168,6 @@ export function getFreebuffModel(id: string): FreebuffModelOption { ) } -function getZonedParts(date: Date, timeZone: string): ZonedDateParts { - const parts = new Intl.DateTimeFormat('en-US', { - timeZone, - year: 'numeric', - month: '2-digit', - day: '2-digit', - hour: '2-digit', - minute: '2-digit', - hourCycle: 'h23', - }).formatToParts(date) - const value = (type: string) => - parts.find((part) => part.type === type)?.value - const year = Number(value('year') ?? 0) - const month = Number(value('month') ?? 1) - const day = Number(value('day') ?? 1) - const hour = Number(value('hour') ?? 0) - const minute = Number(value('minute') ?? 0) - return { - year, - month, - day, - hour, - minute, - } -} - -function addDaysToYmd( - year: number, - month: number, - day: number, - days: number, -): Pick { - const next = new Date(Date.UTC(year, month - 1, day)) - next.setUTCDate(next.getUTCDate() + days) - return { - year: next.getUTCFullYear(), - month: next.getUTCMonth() + 1, - day: next.getUTCDate(), - } -} - -function getUtcForZonedTime( - parts: Pick, - timeZone: string, - hour: number, - minute: number, -): Date { - let guess = new Date( - Date.UTC(parts.year, parts.month - 1, parts.day, hour, minute), - ) - - for (let i = 0; i < 3; i++) { - const actual = getZonedParts(guess, timeZone) - const desiredUtc = Date.UTC( - parts.year, - parts.month - 1, - parts.day, - hour, - minute, - ) - const actualUtc = Date.UTC( - actual.year, - actual.month - 1, - actual.day, - actual.hour, - actual.minute, - ) - guess = new Date(guess.getTime() + (desiredUtc - actualUtc)) - } - - return guess -} - function getNextFreebuffDeploymentStart(now: Date): Date { const easternNow = getZonedParts(now, FREEBUFF_EASTERN_TIMEZONE) const isBeforeTodayOpen = easternNow.hour < 9 diff --git a/common/src/templates/initial-agents-dir/types/tools.ts b/common/src/templates/initial-agents-dir/types/tools.ts index 9cfe1cdf2..cb3882fc0 100644 --- a/common/src/templates/initial-agents-dir/types/tools.ts +++ b/common/src/templates/initial-agents-dir/types/tools.ts @@ -226,10 +226,10 @@ export interface ProposeStrReplaceParams { /** Array of replacements to make. */ replacements: { /** The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation. */ - old: string - /** The string to replace the corresponding old string with. Can be empty to delete. */ - new: string - /** Whether to allow multiple replacements of old string. */ + oldString: string + /** The string to replace the corresponding oldString with. Can be empty to delete. */ + newString: string + /** Whether to allow multiple replacements of oldString. */ allowMultiple?: boolean }[] } @@ -358,10 +358,10 @@ export interface StrReplaceParams { /** Array of replacements to make. */ replacements: { /** The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation. */ - old: string - /** The string to replace the corresponding old string with. Can be empty to delete. */ - new: string - /** Whether to allow multiple replacements of old string. */ + oldString: string + /** The string to replace the corresponding oldString with. Can be empty to delete. */ + newString: string + /** Whether to allow multiple replacements of oldString. */ allowMultiple?: boolean }[] } diff --git a/common/src/tools/params/__tests__/coerce-to-array.test.ts b/common/src/tools/params/__tests__/coerce-to-array.test.ts index ece3e12c4..ccd80ce6b 100644 --- a/common/src/tools/params/__tests__/coerce-to-array.test.ts +++ b/common/src/tools/params/__tests__/coerce-to-array.test.ts @@ -135,8 +135,8 @@ describe('normalizeReplacementAliases', () => { ).toEqual({ old_str: 'before', new_str: 'after', - old: 'before', - new: 'after', + oldString: 'before', + newString: 'after', allowMultiple: true, }) }) @@ -150,22 +150,22 @@ describe('normalizeReplacementAliases', () => { ).toEqual({ old_string: 'before', new_string: 'after', - old: 'before', - new: 'after', + oldString: 'before', + newString: 'after', }) }) it('does not overwrite documented replacement keys', () => { expect( normalizeReplacementAliases({ - old: 'before', - new: 'after', + oldString: 'before', + newString: 'after', old_str: 'ignored', new_str: 'ignored', }), ).toEqual({ - old: 'before', - new: 'after', + oldString: 'before', + newString: 'after', old_str: 'ignored', new_str: 'ignored', }) diff --git a/common/src/tools/params/tool/propose-str-replace.ts b/common/src/tools/params/tool/propose-str-replace.ts index d4d774747..ab86885d7 100644 --- a/common/src/tools/params/tool/propose-str-replace.ts +++ b/common/src/tools/params/tool/propose-str-replace.ts @@ -38,27 +38,27 @@ const inputSchema = z .preprocess( normalizeReplacementAliases, z.object({ - old: z + oldString: z .string() - .min(1, 'Old cannot be empty') + .min(1, 'oldString cannot be empty') .describe( `The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation.`, ), - new: z + newString: z .string() .describe( - `The string to replace the corresponding old string with. Can be empty to delete.`, + `The string to replace the corresponding oldString with. Can be empty to delete.`, ), allowMultiple: z .boolean() .optional() .default(false) .describe( - 'Whether to allow multiple replacements of old string.', + 'Whether to allow multiple replacements of oldString.', ), }), ) - .describe('Pair of old and new strings.'), + .describe('Pair of oldString and newString values.'), ) .min(1, 'Replacements cannot be empty'), ) @@ -79,10 +79,13 @@ ${$getNativeToolCallExampleString({ input: { path: 'path/to/file', replacements: [ - { old: 'This is the old string', new: 'This is the new string' }, { - old: '\nfoo:', - new: '\nbar:', + oldString: 'This is the old string', + newString: 'This is the new string', + }, + { + oldString: '\nfoo:', + newString: '\nbar:', allowMultiple: true, }, ], diff --git a/common/src/tools/params/tool/str-replace.ts b/common/src/tools/params/tool/str-replace.ts index 60350a627..1c697913c 100644 --- a/common/src/tools/params/tool/str-replace.ts +++ b/common/src/tools/params/tool/str-replace.ts @@ -13,7 +13,6 @@ export const updateFileResultSchema = z.union([ z.object({ file: z.string(), message: z.string(), - unifiedDiff: z.string(), }), z.object({ file: z.string(), @@ -39,27 +38,27 @@ const inputSchema = z .preprocess( normalizeReplacementAliases, z.object({ - old: z + oldString: z .string() - .min(1, 'Old cannot be empty') + .min(1, 'oldString cannot be empty') .describe( `The string to replace. This must be an *exact match* of the string you want to replace, including whitespace and punctuation.`, ), - new: z + newString: z .string() .describe( - `The string to replace the corresponding old string with. Can be empty to delete.`, + `The string to replace the corresponding oldString with. Can be empty to delete.`, ), allowMultiple: z .boolean() .optional() .default(false) .describe( - 'Whether to allow multiple replacements of old string.', + 'Whether to allow multiple replacements of oldString.', ), }), ) - .describe('Pair of old and new strings.'), + .describe('Pair of oldString and newString values.'), ) .min(1, 'Replacements cannot be empty'), ) @@ -79,14 +78,18 @@ ${$getNativeToolCallExampleString({ input: { path: 'path/to/file', replacements: [ - { old: 'This is the old string', new: 'This is the new string' }, { - old: '\n\t\t// @codebuff delete this log line please\n\t\tconsole.log("Hello, world!");\n', - new: '\n', + oldString: 'This is the old string', + newString: 'This is the new string', }, { - old: '\nfoo:', - new: '\nbar:', + oldString: + '\n\t\t// @codebuff delete this log line please\n\t\tconsole.log("Hello, world!");\n', + newString: '\n', + }, + { + oldString: '\nfoo:', + newString: '\nbar:', allowMultiple: true, }, ], diff --git a/common/src/tools/params/utils.ts b/common/src/tools/params/utils.ts index 870d7c76c..9b275aa8c 100644 --- a/common/src/tools/params/utils.ts +++ b/common/src/tools/params/utils.ts @@ -43,8 +43,8 @@ export function normalizeReplacementAliases(val: unknown): unknown { const replacement = { ...(val as Record) } for (const [target, aliases] of [ - ['old', ['old_str', 'old_string']], - ['new', ['new_str', 'new_string']], + ['oldString', ['old', 'old_str', 'old_string']], + ['newString', ['new', 'new_str', 'new_string']], ] as const) { if (replacement[target] !== undefined) { continue diff --git a/common/src/types/freebuff-session.ts b/common/src/types/freebuff-session.ts index 6f44d202b..8d4eebd36 100644 --- a/common/src/types/freebuff-session.ts +++ b/common/src/types/freebuff-session.ts @@ -10,13 +10,18 @@ * Usage counter surfaced to the CLI so the waiting-room UI can render * "N of M sessions used" alongside queue/active state. Present when the * joined model consumes premium Freebuff sessions. `recentCount` is the - * rounded session units inside `windowHours` at the time the response was - * produced — see also the standalone `rate_limited` status for the reject - * path. + * rounded session units since the last midnight Pacific reset at the time + * the response was produced — see also the standalone `rate_limited` status + * for the reject path. */ export interface FreebuffSessionRateLimit { model: string limit: number + period: 'pacific_day' + resetTimeZone: string + resetAt: string + /** Deprecated wire field kept for older clients. Premium usage now resets + * at midnight Pacific time rather than using a rolling window. */ windowHours: number recentCount: number } @@ -63,7 +68,7 @@ export type FreebuffSessionServerResponse = * produces `none`). */ queueDepthByModel?: Record /** Current quota snapshots for premium models, keyed by model id. Lets - * the picker show rolling premium-session usage before the user commits + * the picker show today's premium-session usage before the user commits * to a queue. */ rateLimitsByModel?: FreebuffSessionRateLimitByModel } @@ -159,22 +164,23 @@ export type FreebuffSessionServerResponse = status: 'banned' } | { - /** User has used up their shared premium-session quota in the rolling - * window. Returned from POST /session before the user is placed in the - * queue. `retryAfterMs` is the time until enough session units fall out - * of the window to open one quota slot — clients should show the user - * when they can try again. Terminal for the CLI's current poll session; - * the user can exit and come back later. */ + /** User has used up their shared premium-session quota for the current + * Pacific day. Returned from POST /session before the user is placed in + * the queue. `retryAfterMs` is the time until the next midnight Pacific + * reset. Terminal for the CLI's current poll session; the user can exit + * and come back later. */ status: 'rate_limited' /** The freebuff model the user tried to join. */ model: string - /** Max premium session units permitted per window (e.g. 5). */ + /** Max premium session units permitted per Pacific day (e.g. 5). */ limit: number - /** Rolling window size in hours (e.g. 20). */ + period: 'pacific_day' + resetTimeZone: string + resetAt: string + /** Deprecated wire field kept for older clients. */ windowHours: number - /** Premium session units inside the window at check time — will be ≥ limit. */ + /** Premium session units since today's Pacific reset — will be ≥ limit. */ recentCount: number - /** Milliseconds from now until the oldest admission in the window - * exits and the user regains one quota slot. */ + /** Milliseconds from now until the next Pacific midnight reset. */ retryAfterMs: number } diff --git a/common/src/util/__tests__/zoned-time.test.ts b/common/src/util/__tests__/zoned-time.test.ts new file mode 100644 index 000000000..84a0233bd --- /dev/null +++ b/common/src/util/__tests__/zoned-time.test.ts @@ -0,0 +1,35 @@ +import { describe, expect, test } from 'bun:test' + +import { getZonedDayBounds } from '../zoned-time' + +describe('getZonedDayBounds', () => { + test('returns the current Pacific day bounds on a normal day', () => { + const bounds = getZonedDayBounds( + new Date('2026-04-17T16:00:00Z'), + 'America/Los_Angeles', + ) + + expect(bounds.startsAt.toISOString()).toBe('2026-04-17T07:00:00.000Z') + expect(bounds.resetsAt.toISOString()).toBe('2026-04-18T07:00:00.000Z') + }) + + test('handles the shorter spring-forward Pacific day', () => { + const bounds = getZonedDayBounds( + new Date('2026-03-08T09:00:00Z'), + 'America/Los_Angeles', + ) + + expect(bounds.startsAt.toISOString()).toBe('2026-03-08T08:00:00.000Z') + expect(bounds.resetsAt.toISOString()).toBe('2026-03-09T07:00:00.000Z') + }) + + test('handles the longer fall-back Pacific day', () => { + const bounds = getZonedDayBounds( + new Date('2026-11-01T09:00:00Z'), + 'America/Los_Angeles', + ) + + expect(bounds.startsAt.toISOString()).toBe('2026-11-01T07:00:00.000Z') + expect(bounds.resetsAt.toISOString()).toBe('2026-11-02T08:00:00.000Z') + }) +}) diff --git a/common/src/util/zoned-time.ts b/common/src/util/zoned-time.ts new file mode 100644 index 000000000..36e13387f --- /dev/null +++ b/common/src/util/zoned-time.ts @@ -0,0 +1,98 @@ +export interface ZonedDateParts { + year: number + month: number + day: number + hour: number + minute: number +} + +export function getZonedParts(date: Date, timeZone: string): ZonedDateParts { + const parts = new Intl.DateTimeFormat('en-US', { + timeZone, + year: 'numeric', + month: '2-digit', + day: '2-digit', + hour: '2-digit', + minute: '2-digit', + hourCycle: 'h23', + }).formatToParts(date) + + const get = (type: string) => { + const value = parts.find((part) => part.type === type)?.value + if (!value) throw new Error(`Missing ${type} in ${timeZone} date parts`) + return Number(value) + } + + return { + year: get('year'), + month: get('month'), + day: get('day'), + hour: get('hour'), + minute: get('minute'), + } +} + +export function addDaysToYmd( + year: number, + month: number, + day: number, + days: number, +): Pick { + const next = new Date(Date.UTC(year, month - 1, day)) + next.setUTCDate(next.getUTCDate() + days) + return { + year: next.getUTCFullYear(), + month: next.getUTCMonth() + 1, + day: next.getUTCDate(), + } +} + +export function getUtcForZonedTime( + parts: Pick, + timeZone: string, + hour: number, + minute: number, +): Date { + let guess = new Date( + Date.UTC(parts.year, parts.month - 1, parts.day, hour, minute), + ) + + for (let i = 0; i < 3; i++) { + const actual = getZonedParts(guess, timeZone) + const desiredUtc = Date.UTC( + parts.year, + parts.month - 1, + parts.day, + hour, + minute, + ) + const actualUtc = Date.UTC( + actual.year, + actual.month - 1, + actual.day, + actual.hour, + actual.minute, + ) + guess = new Date(guess.getTime() + (desiredUtc - actualUtc)) + } + + return guess +} + +export function getZonedDayBounds( + now: Date, + timeZone: string, +): { startsAt: Date; resetsAt: Date } { + const nowParts = getZonedParts(now, timeZone) + const today = { + year: nowParts.year, + month: nowParts.month, + day: nowParts.day, + } + const tomorrow = addDaysToYmd(today.year, today.month, today.day, 1) + + return { + startsAt: getUtcForZonedTime(today, timeZone, 0, 0), + resetsAt: getUtcForZonedTime(tomorrow, timeZone, 0, 0), + } +} diff --git a/docs/freebuff-waiting-room.md b/docs/freebuff-waiting-room.md index 9ba7354ec..a4a74468b 100644 --- a/docs/freebuff-waiting-room.md +++ b/docs/freebuff-waiting-room.md @@ -162,6 +162,10 @@ The final tick result carries a `queueDepthByModel` map and a single `skipped` r | `FREEBUFF_SESSION_LENGTH_MS` | env | 3_600_000 | Session lifetime | | `SESSION_GRACE_MS` | `web/src/server/free-session/config.ts` | 1_800_000 | Drain window after expiry — gate still admits requests so an in-flight agent can finish, but the CLI is expected to block new prompts. Hard cutoff at `expires_at + grace`. | +### Premium Session Quota + +DeepSeek, Kimi, and legacy GLM share a per-user premium quota. The server counts `free_session_admit` rows from the last midnight in `America/Los_Angeles`; when the user reaches `FREEBUFF_PREMIUM_SESSION_LIMIT`, the next premium `POST /session` is rejected until the next Pacific midnight reset. MiniMax remains unlimited. + ## HTTP API All endpoints authenticate via the standard `Authorization: Bearer ` or `x-codebuff-api-key` header. diff --git a/packages/agent-runtime/src/__tests__/process-str-replace.test.ts b/packages/agent-runtime/src/__tests__/process-str-replace.test.ts index aa8392e25..b7e7fd495 100644 --- a/packages/agent-runtime/src/__tests__/process-str-replace.test.ts +++ b/packages/agent-runtime/src/__tests__/process-str-replace.test.ts @@ -20,7 +20,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -41,7 +43,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -61,7 +65,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -80,7 +86,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -95,7 +103,9 @@ describe('processStrReplace', () => { it('should return error if file content is null and oldStr is not empty', async () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: 'old', new: 'new', allowMultiple: false }], + replacements: [ + { oldString: 'old', newString: 'new', allowMultiple: false }, + ], initialContentPromise: Promise.resolve(null), logger, }) @@ -110,7 +120,7 @@ describe('processStrReplace', () => { it('should return error if oldStr is empty and file exists', async () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: '', new: 'new', allowMultiple: false }], + replacements: [{ oldString: '', newString: 'new', allowMultiple: false }], initialContentPromise: Promise.resolve('content'), logger, }) @@ -129,7 +139,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -150,7 +162,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -169,7 +183,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -191,7 +207,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -208,9 +226,21 @@ describe('processStrReplace', () => { it('should continue processing other replacements even if one fails', async () => { const initialContent = 'const x = 1;\nconst y = 2;\nconst z = 3;\n' const replacements = [ - { old: 'const x = 1;', new: 'const x = 10;', allowMultiple: false }, // This exists - { old: 'const w = 4;', new: 'const w = 40;', allowMultiple: false }, // This doesn't exist - { old: 'const z = 3;', new: 'const z = 30;', allowMultiple: false }, // This also exists + { + oldString: 'const x = 1;', + newString: 'const x = 10;', + allowMultiple: false, + }, // This exists + { + oldString: 'const w = 4;', + newString: 'const w = 40;', + allowMultiple: false, + }, // This doesn't exist + { + oldString: 'const z = 3;', + newString: 'const z = 30;', + allowMultiple: false, + }, // This also exists ] const result = await processStrReplace({ @@ -242,7 +272,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -262,7 +294,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -281,7 +315,9 @@ describe('processStrReplace', () => { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -296,9 +332,9 @@ describe('processStrReplace', () => { it('should handle mixed allowMultiple settings in multiple replacements', async () => { const initialContent = 'foo bar foo\nbaz baz baz\nqux qux' const replacements = [ - { old: 'foo', new: 'FOO', allowMultiple: true }, // Replace all 'foo' - { old: 'baz', new: 'BAZ', allowMultiple: false }, // Should error on multiple 'baz' - { old: 'qux qux', new: 'QUX', allowMultiple: false }, // Single occurrence, should work + { oldString: 'foo', newString: 'FOO', allowMultiple: true }, // Replace all 'foo' + { oldString: 'baz', newString: 'BAZ', allowMultiple: false }, // Should error on multiple 'baz' + { oldString: 'qux qux', newString: 'QUX', allowMultiple: false }, // Single occurrence, should work ] const result = await processStrReplace({ @@ -335,7 +371,9 @@ function test3() { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -359,7 +397,9 @@ function test3() { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -383,7 +423,9 @@ function test3() { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -403,7 +445,9 @@ function test3() { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: true }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: true }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) @@ -422,13 +466,13 @@ function test3() { const initialContent = 'line 1\nline 2\nline 3\n' const replacements = [ { - old: 'line 2\n', - new: 'this is a new line\n', + oldString: 'line 2\n', + newString: 'this is a new line\n', allowMultiple: false, }, { - old: 'line 3\n', - new: 'new line 3\n', + oldString: 'line 3\n', + newString: 'new line 3\n', allowMultiple: false, }, ] @@ -454,7 +498,9 @@ function test3() { const result = await processStrReplace({ path: 'test.ts', - replacements: [{ old: oldStr, new: newStr, allowMultiple: false }], + replacements: [ + { oldString: oldStr, newString: newStr, allowMultiple: false }, + ], initialContentPromise: Promise.resolve(initialContent), logger, }) diff --git a/packages/agent-runtime/src/__tests__/propose-tools.test.ts b/packages/agent-runtime/src/__tests__/propose-tools.test.ts index 84ceafb07..55ae16f4d 100644 --- a/packages/agent-runtime/src/__tests__/propose-tools.test.ts +++ b/packages/agent-runtime/src/__tests__/propose-tools.test.ts @@ -1,10 +1,7 @@ import { TEST_USER_ID } from '@codebuff/common/old-constants' import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime' import { getInitialSessionState } from '@codebuff/common/types/session-state' -import { - assistantMessage, - userMessage, -} from '@codebuff/common/util/messages' +import { assistantMessage, userMessage } from '@codebuff/common/util/messages' import { afterEach, beforeEach, @@ -51,7 +48,9 @@ describe('propose_str_replace and propose_write_file tools', () => { let mockTemplate: AgentTemplate let mockAgentState: AgentState let mockParams: ParamsOf - let executeToolCallSpy: ReturnType> + let executeToolCallSpy: ReturnType< + typeof spyOn + > let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps // Mock file system - maps file paths to their contents @@ -59,7 +58,8 @@ describe('propose_str_replace and propose_write_file tools', () => { beforeEach(() => { // Reset mock file system - mockFiles['src/utils.ts'] = `export function add(a: number, b: number): number { + mockFiles['src/utils.ts'] = + `export function add(a: number, b: number): number { return a + b; } @@ -87,18 +87,27 @@ console.log(add(1, 2)); if (toolName === 'propose_str_replace') { const { path, replacements } = input as { path: string - replacements: Array<{ old: string; new: string; allowMultiple: boolean }> + replacements: Array<{ + oldString: string + newString: string + allowMultiple: boolean + }> } - + // Get current content (from proposed state or mock files) let content = mockFiles[path] ?? null - + if (content === null) { const errorResult: ToolMessage = { role: 'tool', toolName: 'propose_str_replace', toolCallId: `${toolName}-call-id`, - content: [{ type: 'json', value: { file: path, errorMessage: `File not found: ${path}` } }], + content: [ + { + type: 'json', + value: { file: path, errorMessage: `File not found: ${path}` }, + }, + ], } toolResults.push(errorResult) agentState.messageHistory.push(errorResult) @@ -108,14 +117,22 @@ console.log(add(1, 2)); // Apply replacements const errors: string[] = [] for (const replacement of replacements) { - if (!content.includes(replacement.old)) { - errors.push(`String not found: "${replacement.old.slice(0, 50)}..."`) + if (!content.includes(replacement.oldString)) { + errors.push( + `String not found: "${replacement.oldString.slice(0, 50)}..."`, + ) continue } if (replacement.allowMultiple) { - content = content.replaceAll(replacement.old, replacement.new) + content = content.replaceAll( + replacement.oldString, + replacement.newString, + ) } else { - content = content.replace(replacement.old, replacement.new) + content = content.replace( + replacement.oldString, + replacement.newString, + ) } } @@ -124,7 +141,12 @@ console.log(add(1, 2)); role: 'tool', toolName: 'propose_str_replace', toolCallId: `${toolName}-call-id`, - content: [{ type: 'json', value: { file: path, errorMessage: errors.join('; ') } }], + content: [ + { + type: 'json', + value: { file: path, errorMessage: errors.join('; ') }, + }, + ], } toolResults.push(errorResult) agentState.messageHistory.push(errorResult) @@ -134,7 +156,7 @@ console.log(add(1, 2)); // Generate unified diff const originalContent = mockFiles[path]! const diff = generateSimpleDiff(path, originalContent, content) - + // Store proposed content for future calls mockFiles[path] = content @@ -142,14 +164,16 @@ console.log(add(1, 2)); role: 'tool', toolName: 'propose_str_replace', toolCallId: `${toolName}-call-id`, - content: [{ - type: 'json', - value: { - file: path, - message: 'Proposed string replacements', - unifiedDiff: diff, + content: [ + { + type: 'json', + value: { + file: path, + message: 'Proposed string replacements', + unifiedDiff: diff, + }, }, - }], + ], } toolResults.push(successResult) agentState.messageHistory.push(successResult) @@ -159,13 +183,13 @@ console.log(add(1, 2)); instructions: string content: string } - + const originalContent = mockFiles[path] ?? '' const isNewFile = !(path in mockFiles) - + // Generate unified diff const diff = generateSimpleDiff(path, originalContent, newContent) - + // Store proposed content mockFiles[path] = newContent @@ -173,14 +197,18 @@ console.log(add(1, 2)); role: 'tool', toolName: 'propose_write_file', toolCallId: `${toolName}-call-id`, - content: [{ - type: 'json', - value: { - file: path, - message: isNewFile ? `Proposed new file ${path}` : `Proposed changes to ${path}`, - unifiedDiff: diff, + content: [ + { + type: 'json', + value: { + file: path, + message: isNewFile + ? `Proposed new file ${path}` + : `Proposed changes to ${path}`, + unifiedDiff: diff, + }, }, - }], + ], } toolResults.push(successResult) agentState.messageHistory.push(successResult) @@ -201,7 +229,8 @@ console.log(add(1, 2)); // Mock crypto.randomUUID spyOn(crypto, 'randomUUID').mockImplementation( - () => 'mock-uuid-0000-0000-0000-000000000000' as `${string}-${string}-${string}-${string}-${string}`, + () => + 'mock-uuid-0000-0000-0000-000000000000' as `${string}-${string}-${string}-${string}-${string}`, ) // Create mock template for implementor agent @@ -215,10 +244,16 @@ console.log(add(1, 2)); includeMessageHistory: true, inheritParentSystemPrompt: false, mcpServers: {}, - toolNames: ['propose_str_replace', 'propose_write_file', 'set_output', 'end_turn'], + toolNames: [ + 'propose_str_replace', + 'propose_write_file', + 'set_output', + 'end_turn', + ], spawnableAgents: [], systemPrompt: 'You are a code implementor that proposes changes.', - instructionsPrompt: 'Implement the requested changes using propose_str_replace or propose_write_file.', + instructionsPrompt: + 'Implement the requested changes using propose_str_replace or propose_write_file.', stepPrompt: '', handleSteps: undefined, } as AgentTemplate @@ -228,7 +263,8 @@ console.log(add(1, 2)); mockAgentState = { ...sessionState.mainAgentState, agentId: 'test-implementor-id', - runId: 'test-run-id' as `${string}-${string}-${string}-${string}-${string}`, + runId: + 'test-run-id' as `${string}-${string}-${string}-${string}-${string}`, messageHistory: [ userMessage('Add a multiply function to src/utils.ts'), assistantMessage('I will implement the changes.'), @@ -281,23 +317,29 @@ console.log(add(1, 2)); toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'export function subtract(a: number, b: number): number {\n return a - b;\n}', - new: `export function subtract(a: number, b: number): number { + replacements: [ + { + oldString: + 'export function subtract(a: number, b: number): number {\n return a - b;\n}', + newString: `export function subtract(a: number, b: number): number { return a - b; } export function multiply(a: number, b: number): number { return a * b; }`, - allowMultiple: false, - }], + allowMultiple: false, + }, + ], }, } toolResultsCapture.push(step.toolResult) - + const firstResult = step.toolResult?.[0] - const unifiedDiff = firstResult?.type === 'json' ? (firstResult.value as { unifiedDiff?: string })?.unifiedDiff : undefined + const unifiedDiff = + firstResult?.type === 'json' + ? (firstResult.value as { unifiedDiff?: string })?.unifiedDiff + : undefined yield { toolName: 'set_output', input: { @@ -325,9 +367,14 @@ export function multiply(a: number, b: number): number { const toolResult = toolResultsCapture[0] expect(toolResult).toBeDefined() expect(toolResult[0].type).toBe('json') - const jsonResult = toolResult[0] as { type: 'json'; value: { file: string; unifiedDiff: string } } + const jsonResult = toolResult[0] as { + type: 'json' + value: { file: string; unifiedDiff: string } + } expect(jsonResult.value.file).toBe('src/utils.ts') - expect(jsonResult.value.unifiedDiff).toContain('+export function multiply') + expect(jsonResult.value.unifiedDiff).toContain( + '+export function multiply', + ) expect(jsonResult.value.unifiedDiff).toContain('return a * b') }) @@ -339,11 +386,13 @@ export function multiply(a: number, b: number): number { toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'nonexistent string that does not exist in the file', - new: 'replacement', - allowMultiple: false, - }], + replacements: [ + { + oldString: 'nonexistent string that does not exist in the file', + newString: 'replacement', + allowMultiple: false, + }, + ], }, } toolResultsCapture.push(step.toolResult) @@ -356,7 +405,10 @@ export function multiply(a: number, b: number): number { expect(toolResultsCapture).toHaveLength(1) const toolResult = toolResultsCapture[0] - const jsonResult = toolResult[0] as { type: 'json'; value: { errorMessage: string } } + const jsonResult = toolResult[0] as { + type: 'json' + value: { errorMessage: string } + } expect(jsonResult.value.errorMessage).toContain('String not found') }) @@ -369,11 +421,13 @@ export function multiply(a: number, b: number): number { toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'return a + b;', - new: 'return a + b; // addition', - allowMultiple: false, - }], + replacements: [ + { + oldString: 'return a + b;', + newString: 'return a + b; // addition', + allowMultiple: false, + }, + ], }, } toolResultsCapture.push({ step: 1, result: step1.toolResult }) @@ -383,11 +437,13 @@ export function multiply(a: number, b: number): number { toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'return a - b;', - new: 'return a - b; // subtraction', - allowMultiple: false, - }], + replacements: [ + { + oldString: 'return a - b;', + newString: 'return a - b; // subtraction', + allowMultiple: false, + }, + ], }, } toolResultsCapture.push({ step: 2, result: step2.toolResult }) @@ -400,13 +456,19 @@ export function multiply(a: number, b: number): number { await runProgrammaticStep(mockParams) expect(toolResultsCapture).toHaveLength(2) - + // Both replacements should succeed - const result0 = toolResultsCapture[0].result[0] as { type: 'json'; value: { unifiedDiff: string } } - const result1 = toolResultsCapture[1].result[0] as { type: 'json'; value: { unifiedDiff: string } } + const result0 = toolResultsCapture[0].result[0] as { + type: 'json' + value: { unifiedDiff: string } + } + const result1 = toolResultsCapture[1].result[0] as { + type: 'json' + value: { unifiedDiff: string } + } expect(result0.value.unifiedDiff).toContain('// addition') expect(result1.value.unifiedDiff).toContain('// subtraction') - + // Final file should have both changes expect(mockFiles['src/utils.ts']).toContain('// addition') expect(mockFiles['src/utils.ts']).toContain('// subtraction') @@ -439,10 +501,15 @@ export function multiply(a: number, b: number): number { expect(toolResultsCapture).toHaveLength(1) const toolResult = toolResultsCapture[0] - const jsonResult = toolResult[0] as { type: 'json'; value: { file: string; message: string; unifiedDiff: string } } + const jsonResult = toolResult[0] as { + type: 'json' + value: { file: string; message: string; unifiedDiff: string } + } expect(jsonResult.value.file).toBe('src/multiply.ts') expect(jsonResult.value.message).toContain('new file') - expect(jsonResult.value.unifiedDiff).toContain('+export function multiply') + expect(jsonResult.value.unifiedDiff).toContain( + '+export function multiply', + ) }) it('should propose file edit and return unified diff', async () => { @@ -478,10 +545,15 @@ export function multiply(a: number, b: number): number { expect(toolResultsCapture).toHaveLength(1) const toolResult = toolResultsCapture[0] - const jsonResult = toolResult[0] as { type: 'json'; value: { file: string; message: string; unifiedDiff: string } } + const jsonResult = toolResult[0] as { + type: 'json' + value: { file: string; message: string; unifiedDiff: string } + } expect(jsonResult.value.file).toBe('src/utils.ts') expect(jsonResult.value.message).toContain('changes') - expect(jsonResult.value.unifiedDiff).toContain('+export function multiply') + expect(jsonResult.value.unifiedDiff).toContain( + '+export function multiply', + ) }) }) @@ -501,15 +573,19 @@ export function multiply(a: number, b: number): number { toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'return a + b;', - new: 'return a + b; // first change', - allowMultiple: false, - }], + replacements: [ + { + oldString: 'return a + b;', + newString: 'return a + b; // first change', + allowMultiple: false, + }, + ], }, } const step1First = step1.toolResult?.[0] - const step1HasDiff = step1First?.type === 'json' && !!(step1First.value as { unifiedDiff?: string })?.unifiedDiff + const step1HasDiff = + step1First?.type === 'json' && + !!(step1First.value as { unifiedDiff?: string })?.unifiedDiff receivedToolResults.push({ step: 1, toolResult: step1.toolResult, @@ -521,15 +597,19 @@ export function multiply(a: number, b: number): number { toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'return a - b;', - new: 'return a - b; // second change', - allowMultiple: false, - }], + replacements: [ + { + oldString: 'return a - b;', + newString: 'return a - b; // second change', + allowMultiple: false, + }, + ], }, } const step2First = step2.toolResult?.[0] - const step2HasDiff = step2First?.type === 'json' && !!(step2First.value as { unifiedDiff?: string })?.unifiedDiff + const step2HasDiff = + step2First?.type === 'json' && + !!(step2First.value as { unifiedDiff?: string })?.unifiedDiff receivedToolResults.push({ step: 2, toolResult: step2.toolResult, @@ -546,7 +626,9 @@ export function multiply(a: number, b: number): number { }, } const step3First = step3.toolResult?.[0] - const step3HasDiff = step3First?.type === 'json' && !!(step3First.value as { unifiedDiff?: string })?.unifiedDiff + const step3HasDiff = + step3First?.type === 'json' && + !!(step3First.value as { unifiedDiff?: string })?.unifiedDiff receivedToolResults.push({ step: 3, toolResult: step3.toolResult, @@ -561,31 +643,40 @@ export function multiply(a: number, b: number): number { const result = await runProgrammaticStep(mockParams) expect(result.endTurn).toBe(true) - + // Verify we received tool results for all 3 steps expect(receivedToolResults).toHaveLength(3) - + // Step 1: Should have received tool result with unified diff expect(receivedToolResults[0].step).toBe(1) expect(receivedToolResults[0].toolResult).toBeDefined() expect(receivedToolResults[0].hasUnifiedDiff).toBe(true) - const step1Result = receivedToolResults[0].toolResult[0] as { type: 'json'; value: { file: string; unifiedDiff: string } } + const step1Result = receivedToolResults[0].toolResult[0] as { + type: 'json' + value: { file: string; unifiedDiff: string } + } expect(step1Result.value.file).toBe('src/utils.ts') expect(step1Result.value.unifiedDiff).toContain('first change') - + // Step 2: Should have received tool result with unified diff expect(receivedToolResults[1].step).toBe(2) expect(receivedToolResults[1].toolResult).toBeDefined() expect(receivedToolResults[1].hasUnifiedDiff).toBe(true) - const step2Result = receivedToolResults[1].toolResult[0] as { type: 'json'; value: { file: string; unifiedDiff: string } } + const step2Result = receivedToolResults[1].toolResult[0] as { + type: 'json' + value: { file: string; unifiedDiff: string } + } expect(step2Result.value.file).toBe('src/utils.ts') expect(step2Result.value.unifiedDiff).toContain('second change') - + // Step 3: Should have received tool result with unified diff for new file expect(receivedToolResults[2].step).toBe(3) expect(receivedToolResults[2].toolResult).toBeDefined() expect(receivedToolResults[2].hasUnifiedDiff).toBe(true) - const step3Result = receivedToolResults[2].toolResult[0] as { type: 'json'; value: { file: string; message: string } } + const step3Result = receivedToolResults[2].toolResult[0] as { + type: 'json' + value: { file: string; message: string } + } expect(step3Result.value.file).toBe('src/new-file.ts') expect(step3Result.value.message).toContain('new file') }) @@ -607,20 +698,23 @@ export function multiply(a: number, b: number): number { toolName: 'propose_str_replace', input: { path: 'src/utils.ts', - replacements: [{ - old: 'export function subtract(a: number, b: number): number {\n return a - b;\n}', - new: `export function subtract(a: number, b: number): number { + replacements: [ + { + oldString: + 'export function subtract(a: number, b: number): number {\n return a - b;\n}', + newString: `export function subtract(a: number, b: number): number { return a - b; } export function multiply(a: number, b: number): number { return a * b; }`, - allowMultiple: false, - }], + allowMultiple: false, + }, + ], }, } - + // Capture the tool call and result capturedToolCalls.push({ toolName: 'propose_str_replace', @@ -654,7 +748,7 @@ export function multiply(a: number, b: number): number { expect(result.endTurn).toBe(true) expect(result.agentState.output).toBeDefined() - + const output = result.agentState.output as { toolCalls: any[] toolResults: any[] @@ -668,7 +762,9 @@ export function multiply(a: number, b: number): number { // Verify tool results were captured expect(output.toolResults).toHaveLength(1) expect(output.toolResults[0].file).toBe('src/utils.ts') - expect(output.toolResults[0].unifiedDiff).toContain('+export function multiply') + expect(output.toolResults[0].unifiedDiff).toContain( + '+export function multiply', + ) // Verify unified diffs string was generated expect(output.unifiedDiffs).toContain('--- src/utils.ts ---') @@ -681,25 +777,31 @@ export function multiply(a: number, b: number): number { * Simple diff generator for testing purposes. * In production, the actual handlers use the 'diff' library. */ -function generateSimpleDiff(path: string, oldContent: string, newContent: string): string { +function generateSimpleDiff( + path: string, + oldContent: string, + newContent: string, +): string { const oldLines = oldContent.split('\n') const newLines = newContent.split('\n') - + const diffLines: string[] = [] const maxLen = Math.max(oldLines.length, newLines.length) - + let inChange = false let _changeStart = 0 - + for (let i = 0; i < maxLen; i++) { const oldLine = oldLines[i] const newLine = newLines[i] - + if (oldLine !== newLine) { if (!inChange) { inChange = true _changeStart = i - diffLines.push(`@@ -${i + 1},${oldLines.length - i} +${i + 1},${newLines.length - i} @@`) + diffLines.push( + `@@ -${i + 1},${oldLines.length - i} +${i + 1},${newLines.length - i} @@`, + ) } if (oldLine !== undefined) { diffLines.push(`-${oldLine}`) @@ -711,6 +813,6 @@ function generateSimpleDiff(path: string, oldContent: string, newContent: string diffLines.push(` ${oldLine}`) } } - + return diffLines.join('\n') } diff --git a/packages/agent-runtime/src/__tests__/tool-validation-error.test.ts b/packages/agent-runtime/src/__tests__/tool-validation-error.test.ts index ff75aa44e..ed5cfaa5a 100644 --- a/packages/agent-runtime/src/__tests__/tool-validation-error.test.ts +++ b/packages/agent-runtime/src/__tests__/tool-validation-error.test.ts @@ -144,7 +144,32 @@ describe('tool validation error handling', () => { expect('error' in result).toBe(false) if (!('error' in result)) { expect(result.input.replacements).toEqual([ - { old: 'before', new: 'after', allowMultiple: false }, + { oldString: 'before', newString: 'after', allowMultiple: false }, + ]) + } + }) + + it('should accept old/new aliases for str_replace replacements', () => { + const result = parseRawToolCall({ + rawToolCall: { + toolName: 'str_replace', + toolCallId: 'short-alias-tool-call-id', + input: { + path: 'test.ts', + replacements: [ + { + old: 'before', + new: 'after', + }, + ], + }, + }, + }) + + expect('error' in result).toBe(false) + if (!('error' in result)) { + expect(result.input.replacements).toEqual([ + { oldString: 'before', newString: 'after', allowMultiple: false }, ]) } }) @@ -169,7 +194,7 @@ describe('tool validation error handling', () => { expect('error' in result).toBe(false) if (!('error' in result)) { expect(result.input.replacements).toEqual([ - { old: 'before', new: 'after', allowMultiple: false }, + { oldString: 'before', newString: 'after', allowMultiple: false }, ]) } }) @@ -182,9 +207,9 @@ describe('tool validation error handling', () => { input: { path: 'test.ts', replacements: [ - { old: 'before', new: 'after' }, - { old: 'delete me' }, - { old: 'delete me too' }, + { oldString: 'before', newString: 'after' }, + { oldString: 'delete me' }, + { oldString: 'delete me too' }, ], }, }, @@ -193,10 +218,10 @@ describe('tool validation error handling', () => { expect('error' in result).toBe(true) if ('error' in result) { expect(result.error).toContain('Missing required replacement fields:') - expect(result.error).toContain('- replacements[1].new') - expect(result.error).toContain('- replacements[2].new') + expect(result.error).toContain('- replacements[1].newString') + expect(result.error).toContain('- replacements[2].newString') expect(result.error).toContain( - 'If the intent is deletion, set "new": "" explicitly.', + 'If the intent is deletion, set "newString": "" explicitly.', ) expect(result.error).toContain('Raw validation issues:') } diff --git a/packages/agent-runtime/src/process-str-replace.ts b/packages/agent-runtime/src/process-str-replace.ts index 12d25d48d..e836b77fd 100644 --- a/packages/agent-runtime/src/process-str-replace.ts +++ b/packages/agent-runtime/src/process-str-replace.ts @@ -10,7 +10,11 @@ function normalizeLineEndings(params: { str: string }): string { export async function processStrReplace(params: { path: string - replacements: { old: string; new: string; allowMultiple: boolean }[] + replacements: { + oldString: string + newString: string + allowMultiple: boolean + }[] initialContentPromise: Promise logger: Logger }): Promise< @@ -34,12 +38,16 @@ export async function processStrReplace(params: { } } - // Process each old/new string pair + // Process each oldString/newString pair let currentContent = initialContent let messages: string[] = [] const lineEnding = currentContent.includes('\r\n') ? '\r\n' : '\n' - for (const { old: oldStr, new: newStr, allowMultiple } of replacements) { + for (const { + oldString: oldStr, + newString: newStr, + allowMultiple, + } of replacements) { // Regular case: require oldStr for replacements if (!oldStr) { messages.push( diff --git a/packages/agent-runtime/src/tools/tool-executor.ts b/packages/agent-runtime/src/tools/tool-executor.ts index 60993a022..de97e27bf 100644 --- a/packages/agent-runtime/src/tools/tool-executor.ts +++ b/packages/agent-runtime/src/tools/tool-executor.ts @@ -161,7 +161,7 @@ function summarizeMissingReplacementFields( issue.message?.includes('received undefined') && root === 'replacements' && typeof index === 'number' && - (field === 'old' || field === 'new') + (field === 'oldString' || field === 'newString') return isMissingReplacementString ? [`replacements[${index}].${field}`] : [] }) @@ -174,13 +174,13 @@ function summarizeMissingReplacementFields( 'Missing required replacement fields:', ...missingFields.map((field) => `- ${field}`), '', - 'If the intent is deletion, set "new": "" explicitly.', + 'If the intent is deletion, set "newString": "" explicitly.', ].join('\n') } function getToolValidationHint(toolName: string): string | undefined { if (toolName === 'str_replace' || toolName === 'propose_str_replace') { - return 'Expected shape: { "path": string, "replacements": [{ "old": string, "new": string, "allowMultiple"?: boolean }] }.' + return 'Expected shape: { "path": string, "replacements": [{ "oldString": string, "newString": string, "allowMultiple"?: boolean }] }.' } if (toolName === 'write_file' || toolName === 'propose_write_file') { return 'Expected shape: { "path": string, "instructions": string, "content": string }. Quote string values and escape newlines/quotes inside content.' diff --git a/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts b/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts index a61e82703..7b182237b 100644 --- a/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts +++ b/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts @@ -39,7 +39,7 @@ Some text between { "cb_tool_name": "str_replace", "path": "file1.ts", - "replacements": [{"old": "foo", "new": "bar"}] + "replacements": [{"oldString": "foo", "newString": "bar"}] } @@ -56,7 +56,7 @@ Some commentary after` toolName: 'str_replace', input: { path: 'file1.ts', - replacements: [{ old: 'foo', new: 'bar' }], + replacements: [{ oldString: 'foo', newString: 'bar' }], }, }) }) @@ -178,7 +178,7 @@ Some commentary after` '{\n' + ' "cb_tool_name": "str_replace",\n' + ' "path": "test.ts",\n' + - ' "replacements": [{"old": "console.log(\\"hello\\")", "new": "console.log(\'world\')"}]\n' + + ' "replacements": [{"oldString": "console.log(\\"hello\\")", "newString": "console.log(\'world\')"}]\n' + '}\n' + '' @@ -186,10 +186,10 @@ Some commentary after` expect(result).toHaveLength(1) const replacements = result[0].input.replacements as Array<{ - old: string - new: string + oldString: string + newString: string }> - expect(replacements[0].old).toBe('console.log("hello")') + expect(replacements[0].oldString).toBe('console.log("hello")') }) it('should handle tool calls with newlines in content', () => { diff --git a/packages/internal/src/db/schema.ts b/packages/internal/src/db/schema.ts index ee4f32509..79357c2b6 100644 --- a/packages/internal/src/db/schema.ts +++ b/packages/internal/src/db/schema.ts @@ -911,9 +911,9 @@ export const freeSession = pgTable( /** * Audit log of every admission — one row per queued→active transition. Used - * to track shared premium-session usage for Freebuff's 5 sessions / 20h - * allowance. `session_units` starts at 1.0 and may be reduced when users end - * active sessions early. + * to track shared premium-session usage for Freebuff's 5 sessions per Pacific + * day allowance. `session_units` starts at 1.0 and may be reduced when users + * end active sessions early. * * Separate from `free_session` because that table is one-row-per-user (state, * not history); the UPSERT path there would otherwise destroy prior admissions. diff --git a/sdk/src/__tests__/change-file.test.ts b/sdk/src/__tests__/change-file.test.ts new file mode 100644 index 000000000..dff8969c7 --- /dev/null +++ b/sdk/src/__tests__/change-file.test.ts @@ -0,0 +1,96 @@ +import { describe, expect, test } from 'bun:test' + +import { createMockFs } from '@codebuff/common/testing/mocks/filesystem' + +import { changeFile } from '../tools/change-file' + +describe('changeFile', () => { + test('returns a simple success message for string replacements', async () => { + const fs = createMockFs({ + files: { + '/repo/src/file.ts': 'const value = 1\n', + }, + }) + + const result = await changeFile({ + parameters: { + type: 'patch', + path: 'src/file.ts', + content: '@@ -1,1 +1,1 @@\n-const value = 1\n+const value = 2\n', + }, + cwd: '/repo', + fs, + }) + + expect(result).toEqual([ + { + type: 'json', + value: { + file: 'src/file.ts', + message: 'String replace applied successfully.', + }, + }, + ]) + expect(await fs.readFile('/repo/src/file.ts', 'utf-8')).toBe( + 'const value = 2\n', + ) + }) + + test('returns a simple success message for new file writes', async () => { + const fs = createMockFs() + + const result = await changeFile({ + parameters: { + type: 'file', + path: 'src/file.ts', + content: 'const value = 1\n', + }, + cwd: '/repo', + fs, + }) + + expect(result).toEqual([ + { + type: 'json', + value: { + file: 'src/file.ts', + message: 'Created file successfully.', + }, + }, + ]) + expect(await fs.readFile('/repo/src/file.ts', 'utf-8')).toBe( + 'const value = 1\n', + ) + }) + + test('returns a simple success message for overwritten file writes', async () => { + const fs = createMockFs({ + files: { + '/repo/src/file.ts': 'const value = 1\n', + }, + }) + + const result = await changeFile({ + parameters: { + type: 'file', + path: 'src/file.ts', + content: 'const value = 2\n', + }, + cwd: '/repo', + fs, + }) + + expect(result).toEqual([ + { + type: 'json', + value: { + file: 'src/file.ts', + message: 'Overwrote file successfully.', + }, + }, + ]) + expect(await fs.readFile('/repo/src/file.ts', 'utf-8')).toBe( + 'const value = 2\n', + ) + }) +}) diff --git a/sdk/src/tools/change-file.ts b/sdk/src/tools/change-file.ts index da372e7db..ff34cc547 100644 --- a/sdk/src/tools/change-file.ts +++ b/sdk/src/tools/change-file.ts @@ -4,7 +4,6 @@ import { fileExists } from '@codebuff/common/util/file' import { applyPatch } from 'diff' import z from 'zod/v4' - import type { CodebuffToolOutput } from '@codebuff/common/tools/list' import type { CodebuffFileSystem } from '@codebuff/common/types/filesystem' @@ -43,7 +42,6 @@ export async function changeFile(params: { if (containsPathTraversal(fileChange.path)) { throw new Error('file path contains invalid path traversal') } - const lines = fileChange.content.split('\n') const { created, modified, invalid, patchFailed } = await applyChanges({ projectRoot: cwd, @@ -56,16 +54,20 @@ export async function changeFile(params: { for (const file of created) { results.push({ file, - message: 'Created new file', - unifiedDiff: lines.join('\n'), + message: + fileChange.type === 'patch' + ? 'String replace applied successfully.' + : 'Created file successfully.', }) } for (const file of modified) { results.push({ file, - message: 'Updated file', - unifiedDiff: lines.join('\n'), + message: + fileChange.type === 'patch' + ? 'String replace applied successfully.' + : 'Overwrote file successfully.', }) } @@ -73,7 +75,7 @@ export async function changeFile(params: { results.push({ file, errorMessage: `Failed to apply patch.`, - patch: lines.join('\n'), + patch: fileChange.content, }) } diff --git a/web/src/server/free-session/__tests__/public-api.test.ts b/web/src/server/free-session/__tests__/public-api.test.ts index d29c2cb1f..2ac2ad75a 100644 --- a/web/src/server/free-session/__tests__/public-api.test.ts +++ b/web/src/server/free-session/__tests__/public-api.test.ts @@ -23,6 +23,19 @@ import type { InternalSessionRow } from '../types' const SESSION_LEN = 60 * 60 * 1000 const GRACE_MS = 30 * 60 * 1000 const DEFAULT_MODEL = 'minimax/minimax-m2.7' +const DEFAULT_PREMIUM_RESET_AT = '2026-04-18T07:00:00.000Z' + +function expectedRateLimit(model: string, recentCount: number) { + return { + model, + limit: FREEBUFF_PREMIUM_SESSION_LIMIT, + period: 'pacific_day', + resetTimeZone: 'America/Los_Angeles', + resetAt: DEFAULT_PREMIUM_RESET_AT, + windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, + recentCount, + } as const +} interface AdmitRecord { user_id: string @@ -269,12 +282,7 @@ describe('requestSession', () => { expect(state.status).toBe('queued') if (state.status !== 'queued') throw new Error('unreachable') expect(deps.rows.get('u1')?.model).toBe(FREEBUFF_GLM_MODEL_ID) - expect(state.rateLimit).toEqual({ - model: FREEBUFF_GLM_MODEL_ID, - limit: FREEBUFF_PREMIUM_SESSION_LIMIT, - windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, - recentCount: 0, - }) + expect(state.rateLimit).toEqual(expectedRateLimit(FREEBUFF_GLM_MODEL_ID, 0)) }) test('legacy GLM 5.1 active session can be reclaimed outside deployment hours', async () => { @@ -299,12 +307,7 @@ describe('requestSession', () => { expect(state.status).toBe('active') if (state.status !== 'active') throw new Error('unreachable') expect(state.instanceId).not.toBe('inst-pre') - expect(state.rateLimit).toEqual({ - model: FREEBUFF_GLM_MODEL_ID, - limit: FREEBUFF_PREMIUM_SESSION_LIMIT, - windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, - recentCount: 0, - }) + expect(state.rateLimit).toEqual(expectedRateLimit(FREEBUFF_GLM_MODEL_ID, 0)) }) test('queued response includes a per-model depth snapshot for the selector', async () => { @@ -432,9 +435,9 @@ describe('requestSession', () => { expect(s3.status).toBe('active') }) - // Per-user premium session limit (5 units per 20h) — the wire limit is - // hard-coded in public-api.ts, so tests seed the fake admit log directly - // rather than configuring it. + // Per-user premium session limit (5 units per Pacific day) — the wire + // limit is hard-coded in public-api.ts, so tests seed the fake admit log + // directly rather than configuring it. const PREMIUM_MODEL = FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID const KIMI_MODEL = FREEBUFF_KIMI_MODEL_ID const PREMIUM_LIMIT = FREEBUFF_PREMIUM_SESSION_LIMIT @@ -448,7 +451,7 @@ describe('requestSession', () => { deps.admits.push({ user_id: 'u1', model: i === 0 ? KIMI_MODEL : PREMIUM_MODEL, - admitted_at: new Date(now.getTime() - (19 - i) * 60 * 60 * 1000), + admitted_at: new Date(now.getTime() - i * 60 * 60 * 1000), }) } @@ -463,17 +466,38 @@ describe('requestSession', () => { expect(state.limit).toBe(PREMIUM_LIMIT) expect(state.windowHours).toBe(PREMIUM_WINDOW_HOURS) expect(state.recentCount).toBe(PREMIUM_LIMIT) - expect(state.retryAfterMs).toBe(60 * 60 * 1000) + expect(state.retryAfterMs).toBe(15 * 60 * 60 * 1000) expect(deps.rows.has('u1')).toBe(false) }) - test('rate_limited: DeepSeek admit outside 20h window does not count', async () => { - deps._tick(PREMIUM_OPEN_TIME) + test('rate_limited: reset follows Pacific midnight across DST changes', async () => { + deps._tick(new Date('2026-03-08T09:00:00Z')) const now = deps._now() + for (let i = 0; i < PREMIUM_LIMIT; i++) { + deps.admits.push({ + user_id: 'u1', + model: PREMIUM_MODEL, + admitted_at: new Date(now.getTime() - i * 60_000), + }) + } + + const state = await requestSession({ + userId: 'u1', + model: PREMIUM_MODEL, + deps, + }) + + expect(state.status).toBe('rate_limited') + if (state.status !== 'rate_limited') throw new Error('unreachable') + expect(state.retryAfterMs).toBe(22 * 60 * 60 * 1000) + }) + + test('rate_limited: DeepSeek admit before Pacific midnight does not count', async () => { + deps._tick(PREMIUM_OPEN_TIME) deps.admits.push({ user_id: 'u1', model: PREMIUM_MODEL, - admitted_at: new Date(now.getTime() - 21 * 60 * 60 * 1000), + admitted_at: new Date('2026-04-17T06:59:00Z'), }) const state = await requestSession({ @@ -483,21 +507,15 @@ describe('requestSession', () => { }) expect(state.status).toBe('queued') if (state.status !== 'queued') throw new Error('unreachable') - expect(state.rateLimit).toEqual({ - model: PREMIUM_MODEL, - limit: PREMIUM_LIMIT, - windowHours: PREMIUM_WINDOW_HOURS, - recentCount: 0, - }) + expect(state.rateLimit).toEqual(expectedRateLimit(PREMIUM_MODEL, 0)) }) - test('rate_limited: 5th Kimi admit in window blocks the 6th attempt', async () => { + test('rate_limited: 5th Kimi admit today blocks the 6th attempt', async () => { deps._tick(PREMIUM_OPEN_TIME) - // Seed 5 admits inside the 20h window, spaced so we can verify retryAfter - // points at the oldest one sliding off. + // Seed 5 admits inside today's Pacific day. retryAfter points at the + // next Pacific midnight reset, not the oldest admit. const now = deps._now() - // Oldest: 19h ago (still in window). Next 4: 1h, 2h, 3h, 4h ago. - const ages = [19, 4, 3, 2, 1] + const ages = [8, 4, 3, 2, 1] for (const hoursAgo of ages) { deps.admits.push({ user_id: 'u1', @@ -517,8 +535,7 @@ describe('requestSession', () => { expect(state.limit).toBe(PREMIUM_LIMIT) expect(state.windowHours).toBe(PREMIUM_WINDOW_HOURS) expect(state.recentCount).toBe(PREMIUM_LIMIT) - // Oldest admit is 19h ago; slot opens when it hits 20h, i.e. in 1h. - expect(state.retryAfterMs).toBe(60 * 60 * 1000) + expect(state.retryAfterMs).toBe(15 * 60 * 60 * 1000) // Blocked before any row is written — the user doesn't take a queue slot. expect(deps.rows.has('u1')).toBe(false) }) @@ -546,17 +563,13 @@ describe('requestSession', () => { expect(state.windowHours).toBe(PREMIUM_WINDOW_HOURS) }) - test('rate_limited: admits outside the 20h window do not count', async () => { + test("rate_limited: admits before today's Pacific reset do not count", async () => { deps._tick(PREMIUM_OPEN_TIME) - // 5 admits, each just over 20h old → all fall off the window. - const now = deps._now() for (let i = 0; i < 5; i++) { deps.admits.push({ user_id: 'u1', model: PREMIUM_MODEL, - admitted_at: new Date( - now.getTime() - (PREMIUM_WINDOW_HOURS * 60 * 60 * 1000 + 60_000 + i), - ), + admitted_at: new Date(`2026-04-17T06:5${i}:00Z`), }) } const state = await requestSession({ @@ -592,7 +605,7 @@ describe('requestSession', () => { test('queued DeepSeek response carries the current admit count', async () => { deps._tick(PREMIUM_OPEN_TIME) const now = deps._now() - // 2 admits in the window — under the limit so the user still queues. + // 2 admits today — under the limit so the user still queues. deps.admits.push({ user_id: 'u1', model: PREMIUM_MODEL, @@ -609,12 +622,7 @@ describe('requestSession', () => { deps, }) if (state.status !== 'queued') throw new Error('unreachable') - expect(state.rateLimit).toEqual({ - model: PREMIUM_MODEL, - limit: PREMIUM_LIMIT, - windowHours: PREMIUM_WINDOW_HOURS, - recentCount: 2, - }) + expect(state.rateLimit).toEqual(expectedRateLimit(PREMIUM_MODEL, 2)) }) test('rate_limited: fractional premium usage under the cap can start another session', async () => { @@ -623,7 +631,7 @@ describe('requestSession', () => { deps.admits.push({ user_id: 'u1', model: KIMI_MODEL, - admitted_at: new Date(now.getTime() - 19 * 60 * 60 * 1000), + admitted_at: new Date(now.getTime() - 8 * 60 * 60 * 1000), session_units: 0.9, }) for (let i = 0; i < 4; i++) { @@ -655,7 +663,7 @@ describe('requestSession', () => { const now = deps._now() // Seed 5 prior admits (the cap), with the latest one matching the // active row we're about to install. - const ages = [19, 4, 3, 2, 0] + const ages = [8, 4, 3, 2, 0] for (const hoursAgo of ages) { deps.admits.push({ user_id: 'u1', @@ -685,7 +693,7 @@ describe('requestSession', () => { }) expect(state.status).toBe('active') if (state.status !== 'active') throw new Error('unreachable') - // Instance id rotated; quota snapshot still reflects the full window. + // Instance id rotated; quota snapshot still reflects today's usage. expect(state.instanceId).not.toBe('inst-pre') expect(state.rateLimit?.recentCount).toBe(PREMIUM_LIMIT) }) @@ -736,7 +744,7 @@ describe('requestSession', () => { // must be blocked by the quota. deps._tick(PREMIUM_OPEN_TIME) const now = deps._now() - const ages = [19, 4, 3, 2, 1] + const ages = [8, 4, 3, 2, 1] for (const hoursAgo of ages) { deps.admits.push({ user_id: 'u1', @@ -767,7 +775,7 @@ describe('requestSession', () => { test('instant-admit bumps the quota count for the freshly-written admit row', async () => { const admitDeps = makeDeps({ getInstantAdmitCapacity: () => 3 }) admitDeps._tick(PREMIUM_OPEN_TIME) - // 1 existing admit in the window; this new call should instant-admit and + // 1 existing admit today; this new call should instant-admit and // write a second row, so the response's recentCount reflects 2. const now = admitDeps._now() admitDeps.admits.push({ @@ -816,7 +824,7 @@ describe('getSessionState', () => { deps.admits.push({ user_id: 'u1', model: FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, - admitted_at: new Date(now.getTime() - 19 * 60 * 60 * 1000), + admitted_at: new Date(now.getTime() - 60 * 60 * 1000), }) const state = await getSessionState({ userId: 'u1', deps }) @@ -824,12 +832,7 @@ describe('getSessionState', () => { if (state.status !== 'none') throw new Error('unreachable') expect( state.rateLimitsByModel?.[FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID], - ).toEqual({ - model: FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, - limit: FREEBUFF_PREMIUM_SESSION_LIMIT, - windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, - recentCount: 1, - }) + ).toEqual(expectedRateLimit(FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, 1)) }) test('active session with matching instance id returns active', async () => { @@ -891,12 +894,9 @@ describe('getSessionState', () => { deps, }) if (state.status !== 'active') throw new Error('unreachable') - expect(state.rateLimit).toEqual({ - model: FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, - limit: FREEBUFF_PREMIUM_SESSION_LIMIT, - windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, - recentCount: 1, - }) + expect(state.rateLimit).toEqual( + expectedRateLimit(FREEBUFF_DEEPSEEK_V4_PRO_MODEL_ID, 1), + ) }) test('active session only fetches one shared premium quota snapshot', async () => { diff --git a/web/src/server/free-session/public-api.ts b/web/src/server/free-session/public-api.ts index a1a065abe..59af4db81 100644 --- a/web/src/server/free-session/public-api.ts +++ b/web/src/server/free-session/public-api.ts @@ -4,13 +4,16 @@ import { FREEBUFF_DEPLOYMENT_HOURS_LABEL, FREEBUFF_GEMINI_PRO_MODEL_ID, FREEBUFF_PREMIUM_MODEL_IDS, + FREEBUFF_PREMIUM_SESSION_PERIOD, FREEBUFF_PREMIUM_SESSION_LIMIT, + FREEBUFF_PREMIUM_SESSION_RESET_TIMEZONE, FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, isFreebuffModelAvailable, isFreebuffPremiumModelId, isSupportedFreebuffModelId, resolveSupportedFreebuffModel, } from '@codebuff/common/constants/freebuff-models' +import { getZonedDayBounds } from '@codebuff/common/util/zoned-time' import { getInstantAdmitCapacity, @@ -46,34 +49,15 @@ function roundSessionUnits(units: number): number { return Math.round(units * 10) / 10 } -function getRetryAfterMsForPremiumLimit(params: { - admits: Awaited> - totalUnits: number - targetUnits: number - windowMs: number - now: Date -}): number { - let remainingUnits = params.totalUnits - for (const admit of params.admits) { - remainingUnits = roundSessionUnits(remainingUnits - admit.sessionUnits) - if (remainingUnits <= params.targetUnits) { - return Math.max( - 0, - admit.admittedAt.getTime() + params.windowMs - params.now.getTime(), - ) - } - } - return 0 -} - function canStartPremiumSession(snapshot: FreebuffSessionRateLimit): boolean { return snapshot.recentCount < snapshot.limit } +type PremiumQuotaInfo = Omit + interface PremiumQuotaSnapshot { - recentCount: number - admits: Awaited> - windowMs: number + info: PremiumQuotaInfo + resetsAt: Date } async function fetchPremiumQuotaSnapshot( @@ -81,19 +65,28 @@ async function fetchPremiumQuotaSnapshot( deps: SessionDeps, ): Promise { const now = nowOf(deps) - const windowMs = FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS * 60 * 60 * 1000 - const since = new Date(now.getTime() - windowMs) + const premiumDay = getZonedDayBounds( + now, + FREEBUFF_PREMIUM_SESSION_RESET_TIMEZONE, + ) const admits = await deps.listRecentPremiumAdmits({ userId, - since, + since: premiumDay.startsAt, models: FREEBUFF_PREMIUM_MODEL_IDS, }) + const recentCount = roundSessionUnits( + admits.reduce((sum, admit) => sum + admit.sessionUnits, 0), + ) return { - recentCount: roundSessionUnits( - admits.reduce((sum, admit) => sum + admit.sessionUnits, 0), - ), - admits, - windowMs, + info: { + limit: FREEBUFF_PREMIUM_SESSION_LIMIT, + period: FREEBUFF_PREMIUM_SESSION_PERIOD, + resetTimeZone: FREEBUFF_PREMIUM_SESSION_RESET_TIMEZONE, + resetAt: premiumDay.resetsAt.toISOString(), + windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, + recentCount, + }, + resetsAt: premiumDay.resetsAt, } } @@ -103,9 +96,7 @@ function toRateLimitInfo( ): FreebuffSessionRateLimit { return { model, - limit: FREEBUFF_PREMIUM_SESSION_LIMIT, - windowHours: FREEBUFF_PREMIUM_SESSION_WINDOW_HOURS, - recentCount: snapshot.recentCount, + ...snapshot.info, } } @@ -120,8 +111,7 @@ async function fetchRateLimitSnapshot( ): Promise< | { info: FreebuffSessionRateLimit - admits: Awaited> - windowMs: number + resetsAt: Date } | undefined > { @@ -129,8 +119,7 @@ async function fetchRateLimitSnapshot( const snapshot = await fetchPremiumQuotaSnapshot(userId, deps) return { info: toRateLimitInfo(model, snapshot), - admits: snapshot.admits, - windowMs: snapshot.windowMs, + resetsAt: snapshot.resetsAt, } } @@ -185,7 +174,8 @@ export interface SessionDeps { * bound to a given model. Compared against the model's configured * `instantAdmitCapacity` to decide whether a new joiner skips the queue. */ activeCountForModel: (model: string) => Promise - /** Rate-limit helper: oldest-first premium admissions inside the window. */ + /** Rate-limit helper: oldest-first premium admissions since today's + * Pacific midnight reset. */ listRecentPremiumAdmits: (params: { userId: string models: readonly string[] @@ -271,11 +261,14 @@ export type RequestSessionResult = requestedModel: string } | { - /** User has hit the per-model admission quota in the rolling window. + /** User has hit the per-model admission quota for the current Pacific day. * See `FreebuffSessionServerResponse`'s `rate_limited` variant. */ status: 'rate_limited' model: string limit: number + period: 'pacific_day' + resetTimeZone: string + resetAt: string windowHours: number recentCount: number retryAfterMs: number @@ -328,8 +321,8 @@ export async function requestSession(params: { } // Rate-limit check runs before joinOrTakeOver so heavy users never even - // create a queued row. Premium models share one 20h session-unit pool; - // Minimax falls through unchanged as unlimited. + // create a queued row. Premium models share one daily Pacific-time + // session-unit pool; Minimax falls through unchanged as unlimited. // // Takeover/reclaim exception: a user who already holds a queued or // active+unexpired row on this same model is re-anchoring (CLI restart, @@ -357,19 +350,13 @@ export async function requestSession(params: { if (!isReclaim) { const snapshot = await fetchRateLimitSnapshot(params.userId, model, deps) if (snapshot && !canStartPremiumSession(snapshot.info)) { - const retryAfterMs = getRetryAfterMsForPremiumLimit({ - admits: snapshot.admits, - totalUnits: snapshot.info.recentCount, - targetUnits: snapshot.info.limit, - windowMs: snapshot.windowMs, - now, - }) + const retryAfterMs = Math.max( + 0, + snapshot.resetsAt.getTime() - now.getTime(), + ) return { + ...snapshot.info, status: 'rate_limited', - model, - limit: snapshot.info.limit, - windowHours: snapshot.info.windowHours, - recentCount: snapshot.info.recentCount, retryAfterMs, } }