From 7c3d7a85967fd2e46763e7dd8aef551bf4f3e32a Mon Sep 17 00:00:00 2001
From: Enrico Ros <enrico.ros@gmail.com>
Date: Thu, 26 Sep 2024 12:00:56 -0700
Subject: [PATCH] Proper edit mentions

---
 src/apps/beam/AppBeam.tsx                              | 2 +-
 src/apps/chat/components/layout-bar/ChatBarAltBeam.tsx | 7 ++++---
 src/apps/chat/components/message/ChatMessage.tsx       | 2 +-
 src/common/chat-overlay/ConversationHandler.ts         | 2 +-
 src/modules/beam/store-beam-vanilla.ts                 | 7 +++++--
 5 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/apps/beam/AppBeam.tsx b/src/apps/beam/AppBeam.tsx
index d579c0c94..c9e0aa4a4 100644
--- a/src/apps/beam/AppBeam.tsx
+++ b/src/apps/beam/AppBeam.tsx
@@ -22,7 +22,7 @@ function initTestConversation(): DConversation {
 }
 
 function initTestBeamStore(messages: DMessage[], beamStore: BeamStoreApi = createBeamVanillaStore()): BeamStoreApi {
-  beamStore.getState().open(messages, getChatLLMId(), (content) => alert(content));
+  beamStore.getState().open(messages, getChatLLMId(), false, (content) => alert(content));
   return beamStore;
 }
 
diff --git a/src/apps/chat/components/layout-bar/ChatBarAltBeam.tsx b/src/apps/chat/components/layout-bar/ChatBarAltBeam.tsx
index 0e20c2539..309e87432 100644
--- a/src/apps/chat/components/layout-bar/ChatBarAltBeam.tsx
+++ b/src/apps/chat/components/layout-bar/ChatBarAltBeam.tsx
@@ -24,8 +24,9 @@ export function ChatBarAltBeam(props: {
 
 
   // external beam state
-  const { isScattering, isGatheringAny, requiresConfirmation, setIsMaximized, terminateBeam } = useBeamStore(props.beamStore, useShallow((store) => ({
+  const { isEditMode, isScattering, isGatheringAny, requiresConfirmation, setIsMaximized, terminateBeam } = useBeamStore(props.beamStore, useShallow((store) => ({
     // state
+    isEditMode: store.isEditMode,
     isScattering: store.isScattering,
     isGatheringAny: store.isGatheringAny,
     requiresConfirmation: store.isScattering || store.isGatheringAny || store.raysReady > 0,
@@ -76,9 +77,9 @@ export function ChatBarAltBeam(props: {
               : isScattering ? { animation: `${animationColorBeamScatterINV} 5s infinite, ${animationEnterBelow} 0.6s` }
                 : { fontWeight: 'lg' }
           }>
-          {isGatheringAny ? 'Merging...' : isScattering ? 'Beaming...' : 'Beam'}
+          {isGatheringAny ? 'Merging...' : isScattering ? 'Beaming...' : isEditMode ? 'Beam Edit' : 'Beam'}
         </Box>
-        {(!isGatheringAny && !isScattering) && ' Mode'}
+        {(!isGatheringAny && !isScattering && !isEditMode) && ' Mode'}
       </Typography>
 
       {/* Right Close Icon */}
diff --git a/src/apps/chat/components/message/ChatMessage.tsx b/src/apps/chat/components/message/ChatMessage.tsx
index 9a5ac6638..bbc08a50c 100644
--- a/src/apps/chat/components/message/ChatMessage.tsx
+++ b/src/apps/chat/components/message/ChatMessage.tsx
@@ -925,7 +925,7 @@ export function ChatMessage(props: {
                 ? <>Beam <span style={{ opacity: 0.5 }}>from here</span></>
                 : !props.isBottom
                   ? <>Beam <span style={{ opacity: 0.5 }}>this message</span></>
-                  : <Box sx={{ flexGrow: 1, display: 'flex', justifyContent: 'space-between', gap: 1 }}>Beam<KeyStroke variant='outlined' combo='Ctrl + Shift + B' /></Box>}
+                  : <Box sx={{ flexGrow: 1, display: 'flex', justifyContent: 'space-between', gap: 1 }}>Beam Edit<KeyStroke variant='outlined' combo='Ctrl + Shift + B' /></Box>}
             </MenuItem>
           )}
         </CloseableMenu>
diff --git a/src/common/chat-overlay/ConversationHandler.ts b/src/common/chat-overlay/ConversationHandler.ts
index b49bf2012..403e43072 100644
--- a/src/common/chat-overlay/ConversationHandler.ts
+++ b/src/common/chat-overlay/ConversationHandler.ts
@@ -234,7 +234,7 @@ export class ConversationHandler {
       terminateKeepingSettings();
     };
 
-    beamOpen(viewHistory, getChatLLMId(), onBeamSuccess);
+    beamOpen(viewHistory, getChatLLMId(), !!destReplaceMessageId, onBeamSuccess);
     importMessages.length && beamImportRays(importMessages, getChatLLMId());
   }
 
diff --git a/src/modules/beam/store-beam-vanilla.ts b/src/modules/beam/store-beam-vanilla.ts
index 05252ce44..462d75097 100644
--- a/src/modules/beam/store-beam-vanilla.ts
+++ b/src/modules/beam/store-beam-vanilla.ts
@@ -32,6 +32,7 @@ type BeamSuccessCallback = (messageUpdate: Pick<DMessage, 'fragments' | 'generat
 interface RootStateSlice {
 
   isOpen: boolean;
+  isEditMode: boolean;
   isMaximized: boolean;
   inputHistory: DMessage[] | null;
   inputIssues: string | null;
@@ -43,6 +44,7 @@ interface RootStateSlice {
 const initRootStateSlice = (): RootStateSlice => ({
 
   isOpen: false,
+  isEditMode: false,
   isMaximized: false,
   inputHistory: null,
   inputIssues: null,
@@ -54,7 +56,7 @@ const initRootStateSlice = (): RootStateSlice => ({
 export interface RootStoreSlice extends RootStateSlice {
 
   // lifecycle
-  open: (chatHistory: Readonly<DMessage[]>, initialChatLlmId: DLLMId | null, callback: BeamSuccessCallback) => void;
+  open: (chatHistory: Readonly<DMessage[]>, initialChatLlmId: DLLMId | null, isEditMode: boolean, callback: BeamSuccessCallback) => void;
   terminateKeepingSettings: () => void;
   loadBeamConfig: (preset: BeamConfigSnapshot | null) => void;
 
@@ -70,7 +72,7 @@ const createRootSlice: StateCreator<BeamStore, [], [], RootStoreSlice> = (_set,
   ...initRootStateSlice(),
 
 
-  open: (chatHistory: Readonly<DMessage[]>, initialChatLlmId: DLLMId | null, callback: BeamSuccessCallback) => {
+  open: (chatHistory: Readonly<DMessage[]>, initialChatLlmId: DLLMId | null, isEditMode: boolean, callback: BeamSuccessCallback) => {
     const { isOpen: wasAlreadyOpen, terminateKeepingSettings, loadBeamConfig, hadImportedRays, setRayLlmIds, setCurrentGatherLlmId } = _get();
 
     // reset pending operations
@@ -84,6 +86,7 @@ const createRootSlice: StateCreator<BeamStore, [], [], RootStoreSlice> = (_set,
     _set({
       // input
       isOpen: true,
+      isEditMode,
       inputHistory: isValidHistory ? history : null,
       inputIssues: isValidHistory ? null : 'Invalid conversation history: missing user message',
       inputReady: isValidHistory,