@@ -8,23 +8,42 @@ import { useRooPortal } from "@/components/ui/hooks"
88import { vscode } from "@src/utils/vscode"
99import { Checkpoint } from "./schema"
1010
11- type CheckpointMenuProps = {
11+ type CheckpointMenuBaseProps = {
1212 ts : number
1313 commitHash : string
1414 currentHash ?: string
1515 checkpoint : Checkpoint
1616}
17+ type CheckpointMenuControlledProps = {
18+ open : boolean
19+ onOpenChange : ( open : boolean ) => void
20+ }
21+ type CheckpointMenuUncontrolledProps = {
22+ open ?: undefined
23+ onOpenChange ?: undefined
24+ }
25+ type CheckpointMenuProps = CheckpointMenuBaseProps & ( CheckpointMenuControlledProps | CheckpointMenuUncontrolledProps )
1726
18- export const CheckpointMenu = ( { ts, commitHash, currentHash, checkpoint } : CheckpointMenuProps ) => {
27+ export const CheckpointMenu = ( {
28+ ts,
29+ commitHash,
30+ currentHash,
31+ checkpoint,
32+ open,
33+ onOpenChange,
34+ } : CheckpointMenuProps ) => {
1935 const { t } = useTranslation ( )
20- const [ isOpen , setIsOpen ] = useState ( false )
36+ const [ internalOpen , setInternalOpen ] = useState ( false )
2137 const [ isConfirming , setIsConfirming ] = useState ( false )
2238 const portalContainer = useRooPortal ( "roo-portal" )
2339
2440 const isCurrent = currentHash === commitHash
2541
2642 const previousCommitHash = checkpoint ?. from
2743
44+ const isOpen = open ?? internalOpen
45+ const setOpen = onOpenChange ?? setInternalOpen
46+
2847 const onCheckpointDiff = useCallback ( ( ) => {
2948 vscode . postMessage ( {
3049 type : "checkpointDiff" ,
@@ -34,13 +53,23 @@ export const CheckpointMenu = ({ ts, commitHash, currentHash, checkpoint }: Chec
3453
3554 const onPreview = useCallback ( ( ) => {
3655 vscode . postMessage ( { type : "checkpointRestore" , payload : { ts, commitHash, mode : "preview" } } )
37- setIsOpen ( false )
38- } , [ ts , commitHash ] )
56+ setOpen ( false )
57+ } , [ ts , commitHash , setOpen ] )
3958
4059 const onRestore = useCallback ( ( ) => {
4160 vscode . postMessage ( { type : "checkpointRestore" , payload : { ts, commitHash, mode : "restore" } } )
42- setIsOpen ( false )
43- } , [ ts , commitHash ] )
61+ setOpen ( false )
62+ } , [ ts , commitHash , setOpen ] )
63+
64+ const handleOpenChange = useCallback (
65+ ( open : boolean ) => {
66+ setOpen ( open )
67+ if ( ! open ) {
68+ setIsConfirming ( false )
69+ }
70+ } ,
71+ [ setOpen ] ,
72+ )
4473
4574 return (
4675 < div className = "flex flex-row gap-1" >
@@ -49,15 +78,10 @@ export const CheckpointMenu = ({ ts, commitHash, currentHash, checkpoint }: Chec
4978 < span className = "codicon codicon-diff-single" />
5079 </ Button >
5180 </ StandardTooltip >
52- < Popover
53- open = { isOpen }
54- onOpenChange = { ( open ) => {
55- setIsOpen ( open )
56- setIsConfirming ( false )
57- } } >
81+ < Popover open = { isOpen } onOpenChange = { handleOpenChange } >
5882 < StandardTooltip content = { t ( "chat:checkpoint.menu.restore" ) } >
5983 < PopoverTrigger asChild >
60- < Button variant = "ghost" size = "icon" >
84+ < Button variant = "ghost" size = "icon" aria-label = { t ( "chat:checkpoint.menu.restore" ) } >
6185 < span className = "codicon codicon-history" />
6286 </ Button >
6387 </ PopoverTrigger >
@@ -66,47 +90,58 @@ export const CheckpointMenu = ({ ts, commitHash, currentHash, checkpoint }: Chec
6690 < div className = "flex flex-col gap-2" >
6791 { ! isCurrent && (
6892 < div className = "flex flex-col gap-1 group hover:text-foreground" >
69- < Button variant = "secondary" onClick = { onPreview } >
93+ < Button variant = "secondary" onClick = { onPreview } data-testid = "restore-files-btn" >
7094 { t ( "chat:checkpoint.menu.restoreFiles" ) }
7195 </ Button >
7296 < div className = "text-muted transition-colors group-hover:text-foreground" >
7397 { t ( "chat:checkpoint.menu.restoreFilesDescription" ) }
7498 </ div >
7599 </ div >
76100 ) }
77- < div className = "flex flex-col gap-1 group hover:text-foreground" >
101+ { ! isCurrent && (
78102 < div className = "flex flex-col gap-1 group hover:text-foreground" >
79- { ! isConfirming ? (
80- < Button variant = "secondary" onClick = { ( ) => setIsConfirming ( true ) } >
81- { t ( "chat:checkpoint.menu.restoreFilesAndTask" ) }
82- </ Button >
83- ) : (
84- < >
85- < Button variant = "default" onClick = { onRestore } className = "grow" >
86- < div className = "flex flex-row gap-1" >
87- < CheckIcon />
88- < div > { t ( "chat:checkpoint.menu.confirm" ) } </ div >
89- </ div >
90- </ Button >
91- < Button variant = "secondary" onClick = { ( ) => setIsConfirming ( false ) } >
92- < div className = "flex flex-row gap-1" >
93- < Cross2Icon />
94- < div > { t ( "chat:checkpoint.menu.cancel" ) } </ div >
95- </ div >
103+ < div className = "flex flex-col gap-1 group hover:text-foreground" >
104+ { ! isConfirming ? (
105+ < Button
106+ variant = "secondary"
107+ onClick = { ( ) => setIsConfirming ( true ) }
108+ data-testid = "restore-files-and-task-btn" >
109+ { t ( "chat:checkpoint.menu.restoreFilesAndTask" ) }
96110 </ Button >
97- </ >
98- ) }
99- { isConfirming ? (
100- < div className = "text-destructive font-bold" >
101- { t ( "chat:checkpoint.menu.cannotUndo" ) }
102- </ div >
103- ) : (
104- < div className = "text-muted transition-colors group-hover:text-foreground" >
105- { t ( "chat:checkpoint.menu.restoreFilesAndTaskDescription" ) }
106- </ div >
107- ) }
111+ ) : (
112+ < >
113+ < Button
114+ variant = "default"
115+ onClick = { onRestore }
116+ className = "grow"
117+ data-testid = "confirm-restore-btn" >
118+ < div className = "flex flex-row gap-1" >
119+ < CheckIcon />
120+ < div > { t ( "chat:checkpoint.menu.confirm" ) } </ div >
121+ </ div >
122+ </ Button >
123+ < Button variant = "secondary" onClick = { ( ) => setIsConfirming ( false ) } >
124+ < div className = "flex flex-row gap-1" >
125+ < Cross2Icon />
126+ < div > { t ( "chat:checkpoint.menu.cancel" ) } </ div >
127+ </ div >
128+ </ Button >
129+ </ >
130+ ) }
131+ { isConfirming ? (
132+ < div
133+ data-testid = "checkpoint-confirm-warning"
134+ className = "text-destructive font-bold" >
135+ { t ( "chat:checkpoint.menu.cannotUndo" ) }
136+ </ div >
137+ ) : (
138+ < div className = "text-muted transition-colors group-hover:text-foreground" >
139+ { t ( "chat:checkpoint.menu.restoreFilesAndTaskDescription" ) }
140+ </ div >
141+ ) }
142+ </ div >
108143 </ div >
109- </ div >
144+ ) }
110145 </ div >
111146 </ PopoverContent >
112147 </ Popover >
0 commit comments