Merge pull request #600 from w1ne/agent/kyrall-staged-edit-safety

Stage generated Studio edits safely

Author: w1ne

src/studio/StagedEditSlot.tsx

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2026 Andrii Shylenko and kernelCAD contributors
-import { useCallback } from 'react';
+import { useCallback, useState } from 'react';
 import { Check, X } from 'lucide-react';
 import { useShellStore, shellStore } from './store/useShellStore';
 import { useWorkbench } from './context/WorkbenchContext';
@@ -10,10 +10,9 @@ import type { StagedEdit } from './store/shellStore';
 // populated, renders the intent, a minimal line-by-line diff, and
 // approve/reject buttons. When empty, renders the auto-apply placeholder.
 //
-// Approve writes stagedEdit.toCode through workbench.setCode, which
-// flows through the normal recompute pipeline (AST-edit primacy
-// preserved — the proposed text IS the AST-edited result; we just
-// hand the canonical script to the existing setter).
+// Approve writes stagedEdit.toCode through workbench.setCode only if the
+// editor still matches the staged baseline. That keeps generated edits from
+// overwriting intervening human changes.
 
 function computeLineDiff(from: string, to: string): Array<{ kind: 'context' | 'add' | 'del'; text: string }> {
     // Trivial line diff: walk both, mark non-matching lines as add/del.
@@ -87,15 +86,28 @@ function PlaceholderBody() {
 
 export function StagedEditSlot() {
     const { stagedEdit } = useShellStore();
-    const { setCode } = useWorkbench();
+    const { code, setCode } = useWorkbench();
+    const [staleWarning, setStaleWarning] = useState<{ editId: string; message: string } | null>(null);
+    const visibleStaleWarning =
+        stagedEdit != null && staleWarning?.editId === stagedEdit.id
+            ? staleWarning.message
+            : null;
 
     const handleApprove = useCallback(() => {
         if (stagedEdit == null) return;
+        if (code !== stagedEdit.fromCode) {
+            setStaleWarning({
+                editId: stagedEdit.id,
+                message: 'The editor changed since this edit was staged. Review the current code before applying this proposal.',
+            });
+            return;
+        }
         setCode(stagedEdit.toCode);
         shellStore.clearStagedEdit();
-    }, [stagedEdit, setCode]);
+    }, [code, stagedEdit, setCode]);
 
     const handleReject = useCallback(() => {
+        setStaleWarning(null);
         shellStore.clearStagedEdit();
     }, []);
 
@@ -121,6 +133,14 @@ export function StagedEditSlot() {
                             {stagedEdit.source.kind} · {stagedEdit.source.label}
                         </div>
                     )}
+                    {visibleStaleWarning != null && (
+                        <div
+                            className="rounded border border-amber-800/70 bg-amber-950/40 px-2 py-1 text-[10px] text-amber-200"
+                            data-testid="staged-edit-stale-warning"
+                        >
+                            {visibleStaleWarning}
+                        </div>
+                    )}
                     <div className="flex gap-2">
                         <button
                             type="button"

src/studio/StudioGenerate.tsx

@@ -8,9 +8,10 @@ import { useTextTo3dPreview } from '../funnel/hooks/useTextTo3dPreview';
 import { inAppAgentEnabled } from './agentAvailability';
 import { ConceptResult } from './components/ConceptResult';
 import { useCode } from './context/CodeContext';
-import { useGeometry } from './context/GeometryContext';
 import { useFeatureSelection } from './hooks/useFeatureSelection';
 import { useShellStore, shellStore } from './store/useShellStore';
+import type { AgentRepairWorkflow } from './store/shellStore';
+import type { SelectedFeatureId } from './types';
 
 /** Web-only gate. No hooks here, so the conditional return is safe. */
 export const StudioGenerate: React.FC = () => {
@@ -32,21 +33,28 @@ function stepLabel(e: GenerateEvent): string | null {
     }
 }
 
+interface GenerationReviewSnapshot {
+    readonly fromCode: string;
+    readonly promptText: string;
+    readonly selectedFeatureId: SelectedFeatureId;
+    readonly repairWorkflow: AgentRepairWorkflow | null;
+}
+
 const StudioGenerateInner: React.FC = () => {
     const { phase, events, submit } = useGeneration();
-    const { code, setCode } = useCode();
+    const { code } = useCode();
     const currentCode = code ?? '';
-    const { executeGeometry } = useGeometry();
     const { selectedFeatureId } = useFeatureSelection();
     const { agentDraftPrompt, agentDraftPromptVersion, agentRepairWorkflow } = useShellStore();
     const [editedPrompt, setEditedPrompt] = useState('');
     const [acknowledgedDraftVersion, setAcknowledgedDraftVersion] = useState(-1);
-    // The generationId we've already accepted/rejected — gates the review panel
-    // so an applied/dismissed proposal doesn't reappear.
-    const [resolvedId, setResolvedId] = useState<string | null>(null);
+    // The generationId we've already staged/rejected — gates the review panel
+    // so a resolved proposal doesn't reappear.
+    const [resolution, setResolution] = useState<{ generationId: string; action: 'staged' | 'discarded' } | null>(null);
     // The editor source captured at submit time — the "before" side of the diff
     // (so the diff is stable even though `code` changes once we apply).
     const [baseline, setBaseline] = useState('');
+    const [reviewSnapshot, setReviewSnapshot] = useState<GenerationReviewSnapshot | null>(null);
 
     const prompt =
         agentDraftPrompt !== null && agentDraftPromptVersion !== acknowledgedDraftVersion
@@ -69,7 +77,7 @@ const StudioGenerateInner: React.FC = () => {
     // One operation at a time: the rail is too narrow to narrate two runs.
     const busy = agentBusy || conceptBusy;
     // A finished, not-yet-resolved proposal → show the review (diff + accept/reject).
-    const reviewing = phase.state === 'done' && resolvedId !== phase.generationId;
+    const reviewing = phase.state === 'done' && resolution?.generationId !== phase.generationId;
 
     const steps = useMemo(() => events.map(stepLabel).filter(Boolean) as string[], [events]);
 
@@ -78,12 +86,13 @@ const StudioGenerateInner: React.FC = () => {
         shellStore.setAgentRepairWorkflow({ ...agentRepairWorkflow, state: 'drafted' });
     }, [agentRepairWorkflow, phase.state]);
 
-    const runAgent = (text: string) => {
+    const runAgent = (text: string, snapshot: GenerationReviewSnapshot) => {
         // Edit mode: hand the agent the current model so it iterates instead of
         // replacing. Empty editor → fresh generation.
-        setBaseline(currentCode);
-        if (currentCode.trim()) {
-            void submit(text, currentCode);
+        setBaseline(snapshot.fromCode);
+        setReviewSnapshot(snapshot);
+        if (snapshot.fromCode.trim()) {
+            void submit(text, snapshot.fromCode);
             return;
         }
         void submit(text);
@@ -94,15 +103,22 @@ const StudioGenerateInner: React.FC = () => {
         const trimmed = prompt.trim();
         if (!trimmed || busy) return;
         const agentPrompt = selectedFeatureId === null ? trimmed : `Edit selected target "${selectedFeatureId}": ${trimmed}`;
+        let repairWorkflowForRun = agentRepairWorkflow;
         if (
             agentRepairWorkflow != null &&
             agentRepairWorkflow.state === 'drafted' &&
             agentRepairWorkflow.promptText === trimmed &&
             agentRepairWorkflow.targetId === selectedFeatureId
         ) {
-            shellStore.setAgentRepairWorkflow({ ...agentRepairWorkflow, state: 'running' });
+            repairWorkflowForRun = { ...agentRepairWorkflow, state: 'running' };
+            shellStore.setAgentRepairWorkflow(repairWorkflowForRun);
         }
-        runAgent(agentPrompt);
+        runAgent(agentPrompt, {
+            fromCode: currentCode,
+            promptText: trimmed,
+            selectedFeatureId,
+            repairWorkflow: repairWorkflowForRun,
+        });
     };
 
     const onConcept = () => {
@@ -120,6 +136,12 @@ const StudioGenerateInner: React.FC = () => {
         // review diff still uses the current editor code as its baseline, so
         // nothing is overwritten without the user accepting.
         setBaseline(currentCode);
+        setReviewSnapshot({
+            fromCode: currentCode,
+            promptText: conceptPrompt,
+            selectedFeatureId,
+            repairWorkflow: agentRepairWorkflow,
+        });
         // Read the concept mesh directly from the live preview phase (no mirrored
         // state). A done preview with no Tripo render/fingerprint yields
         // {renderImageUrl:null, proportions:null} — intentional and distinct from
@@ -130,19 +152,34 @@ const StudioGenerateInner: React.FC = () => {
         void submit(conceptPrompt, undefined, mesh);
     };
 
-    const accept = () => {
+    const stageGeneratedEdit = () => {
         if (phase.state !== 'done') return;
-        setCode(phase.artifact.code);
-        void executeGeometry(phase.artifact.code);
-        setResolvedId(phase.generationId);
+        const snapshot = reviewSnapshot ?? {
+            fromCode: currentCode,
+            promptText: prompt.trim(),
+            selectedFeatureId,
+            repairWorkflow: agentRepairWorkflow,
+        };
+        shellStore.proposeStagedEdit({
+            id: `agent:${phase.generationId}`,
+            intent: phase.artifact.title,
+            fromCode: snapshot.fromCode,
+            toCode: phase.artifact.code,
+            source: { kind: 'agent', label: 'Studio Generate' },
+            context: {
+                promptText: snapshot.promptText,
+                selectedFeatureId: snapshot.selectedFeatureId,
+                repairWorkflow: snapshot.repairWorkflow,
+                generationId: phase.generationId,
+            },
+        });
+        setResolution({ generationId: phase.generationId, action: 'staged' });
     };
     const reject = () => {
         if (phase.state !== 'done') return;
-        setResolvedId(phase.generationId);
+        setResolution({ generationId: phase.generationId, action: 'discarded' });
     };
 
-    const isEdit = baseline.trim().length > 0;
-
     return (
         <div className="p-3 flex flex-col gap-2">
             <div className="uppercase tracking-wide text-[10px] text-gray-500">Agent</div>
@@ -222,10 +259,10 @@ const StudioGenerateInner: React.FC = () => {
                     <div className="flex gap-2">
                         <button
                             type="button"
-                            onClick={accept}
+                            onClick={stageGeneratedEdit}
                             className="flex-1 rounded bg-green-600 hover:bg-green-500 text-white px-3 py-1.5 text-[11px] font-medium transition-colors"
                         >
-                            {isEdit ? 'Apply changes' : 'Use this'}
+                            Stage edit
                         </button>
                         <button
                             type="button"
@@ -238,9 +275,14 @@ const StudioGenerateInner: React.FC = () => {
                 </div>
             )}
 
-            {phase.state === 'done' && !reviewing && (
+            {phase.state === 'done' && !reviewing && resolution?.action === 'staged' && (
                 <div className="text-[10px] text-green-500 truncate" aria-live="polite">
-                    ✓ applied — {phase.artifact.title}
+                    ✓ staged for review — {phase.artifact.title}
+                </div>
+            )}
+            {phase.state === 'done' && !reviewing && resolution?.action === 'discarded' && (
+                <div className="text-[10px] text-gray-500 truncate" aria-live="polite">
+                    discarded — {phase.artifact.title}
                 </div>
             )}
             {phase.state === 'error' && (

src/studio/__tests__/StagedEditSlot.test.tsx

@@ -7,14 +7,16 @@ import { StagedEditSlot } from '../StagedEditSlot';
 import { shellStore } from '../store/shellStore';
 
 const setCodeMock = vi.fn();
+let workbenchCode = '';
 
 vi.mock('../context/WorkbenchContext', () => ({
-    useWorkbench: () => ({ setCode: setCodeMock }),
+    useWorkbench: () => ({ code: workbenchCode, setCode: setCodeMock }),
 }));
 
 beforeEach(() => {
     shellStore.reset();
     setCodeMock.mockReset();
+    workbenchCode = '';
 });
 
 afterEach(() => {
@@ -45,10 +47,12 @@ describe('StagedEditSlot', () => {
 
     it('Approve calls workbench.setCode(toCode) and clears the slot', () => {
         const toCode = 'const a = box(10);\nreturn a;';
+        const fromCode = 'return box(10);';
+        workbenchCode = fromCode;
         shellStore.proposeStagedEdit({
             id: 'e2',
             intent: 'demo',
-            fromCode: 'return box(10);',
+            fromCode,
             toCode,
         });
         const { getByTestId } = render(<StagedEditSlot />);
@@ -58,6 +62,23 @@ describe('StagedEditSlot', () => {
         expect(shellStore.getSnapshot().stagedEdit).toBeNull();
     });
 
+    it('Approve blocks stale staged edits when the editor changed since proposal', () => {
+        shellStore.proposeStagedEdit({
+            id: 'e-stale',
+            intent: 'demo',
+            fromCode: 'return box(10);',
+            toCode: 'return box(20);',
+        });
+        workbenchCode = 'return box(15);';
+
+        const { getByTestId } = render(<StagedEditSlot />);
+        fireEvent.click(getByTestId('staged-edit-approve'));
+
+        expect(setCodeMock).not.toHaveBeenCalled();
+        expect(shellStore.getSnapshot().stagedEdit?.id).toBe('e-stale');
+        expect(getByTestId('staged-edit-stale-warning').textContent).toContain('changed since this edit was staged');
+    });
+
     it('Reject leaves the script unchanged and clears the slot', () => {
         shellStore.proposeStagedEdit({
             id: 'e3',