Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 5,443 Bytes
77324b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | import { useEffect, useMemo, useState } from 'react';
import {
Button,
Dialog,
DialogActions,
DialogContent,
DialogTitle,
TextField,
Tooltip,
Typography,
} from '@mui/material';
import BoltOutlinedIcon from '@mui/icons-material/BoltOutlined';
import { useSessionStore } from '@/store/sessionStore';
import { apiFetch } from '@/utils/api';
const DEFAULT_CAP_USD = 5;
function money(value: number | null | undefined): string {
if (value === null || value === undefined) return 'uncapped';
if (value >= 100) return `$${value.toFixed(0)}`;
return `$${value.toFixed(2).replace(/\.00$/, '')}`;
}
export default function YoloControl() {
const { sessions, activeSessionId, updateSessionYolo } = useSessionStore();
const activeSession = useMemo(
() => sessions.find((s) => s.id === activeSessionId) || null,
[sessions, activeSessionId],
);
const [dialogOpen, setDialogOpen] = useState(false);
const [capInput, setCapInput] = useState(String(DEFAULT_CAP_USD));
const [busy, setBusy] = useState(false);
const [error, setError] = useState<string | null>(null);
const enabled = Boolean(activeSession?.autoApprovalEnabled);
const disabled = !activeSessionId || activeSession?.expired || busy;
const remaining = activeSession?.autoApprovalRemainingUsd ?? null;
const cap = activeSession?.autoApprovalCostCapUsd ?? null;
useEffect(() => {
if (!activeSession) return;
setCapInput(String(activeSession.autoApprovalCostCapUsd ?? DEFAULT_CAP_USD));
}, [activeSession?.id, activeSession?.autoApprovalCostCapUsd]); // eslint-disable-line react-hooks/exhaustive-deps
async function patchPolicy(nextEnabled: boolean, nextCap?: number) {
if (!activeSessionId) return null;
setBusy(true);
setError(null);
try {
const body: Record<string, unknown> = { enabled: nextEnabled };
if (nextCap !== undefined) body.cost_cap_usd = nextCap;
const response = await apiFetch(`/api/session/${activeSessionId}/yolo`, {
method: 'PATCH',
body: JSON.stringify(body),
});
if (!response.ok) {
throw new Error(await response.text());
}
const data = await response.json();
updateSessionYolo(activeSessionId, data);
return data;
} catch {
setError('Could not update YOLO settings.');
return null;
} finally {
setBusy(false);
}
}
const handleToggle = async () => {
if (disabled) return;
if (enabled) {
await patchPolicy(false);
return;
}
const nextCap = cap ?? DEFAULT_CAP_USD;
const updated = await patchPolicy(true, nextCap);
if (updated) {
setCapInput(String(updated.cost_cap_usd ?? nextCap));
setDialogOpen(true);
}
};
const handleSaveCap = async () => {
const parsed = Number(capInput);
if (!Number.isFinite(parsed) || parsed < 0) {
setError('Enter a non-negative dollar amount.');
return;
}
const updated = await patchPolicy(true, parsed);
if (updated) setDialogOpen(false);
};
return (
<>
<Tooltip title={enabled ? 'Disable session YOLO auto-approval' : 'Enable session YOLO auto-approval'}>
<span>
<Button
size="small"
variant={enabled ? 'contained' : 'outlined'}
disabled={disabled}
onClick={handleToggle}
startIcon={<BoltOutlinedIcon sx={{ fontSize: 16 }} />}
sx={{
minWidth: { xs: 74, md: 116 },
height: 32,
px: { xs: 1, md: 1.25 },
borderRadius: '8px',
textTransform: 'none',
fontSize: '0.72rem',
whiteSpace: 'nowrap',
bgcolor: enabled ? 'var(--accent-yellow)' : 'transparent',
color: enabled ? '#111' : 'text.secondary',
borderColor: enabled ? 'var(--accent-yellow)' : 'divider',
'&:hover': {
bgcolor: enabled ? 'var(--accent-yellow)' : 'action.hover',
borderColor: 'var(--accent-yellow)',
},
}}
>
{enabled ? `YOLO ${money(remaining)}` : 'YOLO'}
</Button>
</span>
</Tooltip>
<Dialog open={dialogOpen} onClose={() => setDialogOpen(false)} maxWidth="xs" fullWidth>
<DialogTitle sx={{ pb: 1 }}>YOLO Budget</DialogTitle>
<DialogContent sx={{ display: 'flex', flexDirection: 'column', gap: 1.5, pt: 1 }}>
<Typography variant="body2" color="text.secondary">
Auto-approval is active for this session. Scheduled HF jobs still require approval.
</Typography>
<TextField
autoFocus
label="Session cap (USD)"
type="number"
size="small"
value={capInput}
onChange={(e) => setCapInput(e.target.value)}
inputProps={{ min: 0, step: 0.5 }}
error={Boolean(error)}
helperText={error || `Estimated spend: ${money(activeSession?.autoApprovalEstimatedSpendUsd ?? 0)} of ${money(cap)}`}
/>
</DialogContent>
<DialogActions>
<Button onClick={() => setDialogOpen(false)} sx={{ textTransform: 'none' }}>
Close
</Button>
<Button onClick={handleSaveCap} disabled={busy} variant="contained" sx={{ textTransform: 'none' }}>
Save
</Button>
</DialogActions>
</Dialog>
</>
);
}
|