From 5b6c1fa9f97758075bf4027c7d607410be712615 Mon Sep 17 00:00:00 2001
From: Zachary Hamm <zack@systeminit.com>
Date: Thu, 19 Dec 2024 11:51:28 -0600
Subject: [PATCH] feat: auto select components created by mgmt fns

If you create a bunch of components via a management function, (say, via
a template func), we will now replace the selection with them when the
components are created.
---
 .../src/components/ManagementRunPrototype.vue | 22 ---------
 .../ModelingView/ModelingRightClickMenu.vue   | 22 +--------
 app/web/src/store/func/funcs.store.ts         | 35 ++++++++++++++
 app/web/src/store/realtime/realtime_events.ts |  9 ++++
 lib/dal/src/management/mod.rs                 |  7 ++-
 lib/dal/src/management/prototype.rs           | 37 +++++++++++++--
 lib/dal/src/ws_event.rs                       |  5 +-
 lib/sdf-server/src/service/v2/management.rs   | 47 ++++++++++++++-----
 8 files changed, 124 insertions(+), 60 deletions(-)

diff --git a/app/web/src/components/ManagementRunPrototype.vue b/app/web/src/components/ManagementRunPrototype.vue
index dfb7a44f3d..a960ce3b0c 100644
--- a/app/web/src/components/ManagementRunPrototype.vue
+++ b/app/web/src/components/ManagementRunPrototype.vue
@@ -40,7 +40,6 @@
 <script lang="ts" setup>
 import { ref, computed, onMounted, watch } from "vue";
 import clsx from "clsx";
-import { useToast } from "vue-toastification";
 import {
   DropdownMenu,
   DropdownMenuItem,
@@ -69,7 +68,6 @@ import FuncRunTabDropdown from "./FuncRunTabDropdown.vue";
 const funcStore = useFuncStore();
 const viewStore = useViewsStore();
 const router = useRouter();
-const toast = useToast();
 const managementRunsStore = useManagementRunsStore();
 const viewsStore = useViewsStore();
 
@@ -129,27 +127,7 @@ const runPrototype = async (viewId: ViewId) => {
   );
 
   if (result.result.success) {
-    // TODO(WENDY) - for template functions, select the returned componentIds
-    // modelingEventBus.emit("setSelection", componentIds);
-
     lastExecution.value = result.result.data;
-    if (result.result.data.message) {
-      const toastOptions = {
-        pauseOnHover: true,
-        timeout: 5000,
-      };
-      if (result.result.data.status === "ok") {
-        toast.success(
-          `${props.prototype.label}: ${result.result.data.message}`,
-          toastOptions,
-        );
-      } else {
-        toast.warning(
-          `${props.prototype.label}: ${result.result.data.message}`,
-          toastOptions,
-        );
-      }
-    }
   }
 };
 
diff --git a/app/web/src/components/ModelingView/ModelingRightClickMenu.vue b/app/web/src/components/ModelingView/ModelingRightClickMenu.vue
index e52f410206..294714c2fc 100644
--- a/app/web/src/components/ModelingView/ModelingRightClickMenu.vue
+++ b/app/web/src/components/ModelingView/ModelingRightClickMenu.vue
@@ -43,7 +43,6 @@ import { storeToRefs } from "pinia";
 import { computed, ref } from "vue";
 // import plur from "plur";
 import { RouteLocationRaw } from "vue-router";
-import { useToast } from "vue-toastification";
 import { IRect } from "konva/lib/types";
 import { ComponentType } from "@/api/sdf/dal/schema";
 import { useComponentsStore } from "@/store/components.store";
@@ -68,7 +67,6 @@ import {
 
 const contextMenuRef = ref<InstanceType<typeof DropdownMenu>>();
 
-const toast = useToast();
 const changeSetsStore = useChangeSetsStore();
 const componentsStore = useComponentsStore();
 const funcStore = useFuncStore();
@@ -519,29 +517,11 @@ const runManagementFunc = async (prototype: MgmtPrototype) => {
   if (!selectedComponent.value) return;
   if (!viewStore.selectedViewId) return;
 
-  const result = await funcStore.RUN_MGMT_PROTOTYPE(
+  await funcStore.RUN_MGMT_PROTOTYPE(
     prototype.managementPrototypeId,
     selectedComponent.value.def.id,
     viewStore.selectedViewId,
   );
-
-  if (result.result.success && result.result.data.message) {
-    const toastOptions = {
-      pauseOnHover: true,
-      timeout: 5000,
-    };
-    if (result.result.data.status === "ok") {
-      toast.success(
-        `${prototype.label}: ${result.result.data.message}`,
-        toastOptions,
-      );
-    } else {
-      toast.warning(
-        `${prototype.label}: ${result.result.data.message}`,
-        toastOptions,
-      );
-    }
-  }
 };
 
 function triggerCopySelection() {
diff --git a/app/web/src/store/func/funcs.store.ts b/app/web/src/store/func/funcs.store.ts
index a64f944720..5d26a9a488 100644
--- a/app/web/src/store/func/funcs.store.ts
+++ b/app/web/src/store/func/funcs.store.ts
@@ -2,6 +2,7 @@ import * as _ from "lodash-es";
 import { defineStore } from "pinia";
 import { addStoreHooks, ApiRequest, URLPattern } from "@si/vue-lib/pinia";
 import { ulid } from "ulid";
+import { useToast } from "vue-toastification";
 import { Visibility } from "@/api/sdf/dal/visibility";
 import {
   FuncArgument,
@@ -83,6 +84,7 @@ export type MgmtPrototype = {
 export type MgmtPrototypeResult = {
   message?: string;
   status: "ok" | "error";
+  createdComponentIds?: ComponentId[];
 };
 
 export function actionBindingsForVariant(
@@ -207,6 +209,7 @@ export const useFuncStore = () => {
         executingPrompt: false,
         // So we can ignore websocket update originated by this client
         clientUlid: ulid(),
+        pendingManagementRunRequestUlid: undefined as string | undefined,
       }),
       getters: {
         selectedFuncSummary(state): FuncSummary | undefined {
@@ -318,6 +321,9 @@ export const useFuncStore = () => {
               { componentId },
               { viewId },
             ]),
+            optimistic: (requestUlid) => {
+              this.pendingManagementRunRequestUlid = requestUlid;
+            },
           });
         },
 
@@ -940,6 +946,35 @@ export const useFuncStore = () => {
               }
             },
           },
+          {
+            eventType: "ManagementOperationsComplete",
+            callback: (
+              { requestUlid, status, message, createdComponentIds, funcName },
+              _metadata,
+            ) => {
+              if (
+                requestUlid &&
+                requestUlid === this.pendingManagementRunRequestUlid &&
+                createdComponentIds?.length
+              ) {
+                viewStore.setSelectedComponentId(createdComponentIds);
+                this.pendingManagementRunRequestUlid = undefined;
+              }
+
+              const toast = useToast();
+              const toastOptions = {
+                timeout: 5000,
+              };
+              const toastMessage = message?.length
+                ? `Ran ${funcName}: ${message}`
+                : `Ran ${funcName}`;
+              if (status === "ok") {
+                toast.success(toastMessage, toastOptions);
+              } else {
+                toast.warning(toastMessage, toastOptions);
+              }
+            },
+          },
         ]);
         return () => {
           realtimeStore.unsubscribe(this.$id);
diff --git a/app/web/src/store/realtime/realtime_events.ts b/app/web/src/store/realtime/realtime_events.ts
index 41cb073cbe..66a1a9fed9 100644
--- a/app/web/src/store/realtime/realtime_events.ts
+++ b/app/web/src/store/realtime/realtime_events.ts
@@ -215,12 +215,21 @@ export type WsEventPayloadMap = {
         fromComponentId: string;
         toComponentId: string;
       };
+
   ManagementFuncExecuted: {
     managerComponentId: string;
     prototypeId: string;
     funcRunId: FuncRunId;
     changeSetId: string;
   };
+  ManagementOperationsComplete: {
+    requestUlid?: string;
+    funcName: string;
+    status: "ok" | "error";
+    message?: string;
+    createdComponentIds?: ComponentId[];
+  };
+
   ModuleImported: SchemaVariant[];
   WorkspaceImportBeginApprovalProcess: {
     workspacePk: WorkspacePk;
diff --git a/lib/dal/src/management/mod.rs b/lib/dal/src/management/mod.rs
index 848fe37bd3..d3f6e697ad 100644
--- a/lib/dal/src/management/mod.rs
+++ b/lib/dal/src/management/mod.rs
@@ -1010,7 +1010,7 @@ impl<'a> ManagementOperator<'a> {
         Ok(())
     }
 
-    pub async fn operate(&mut self) -> ManagementResult<()> {
+    pub async fn operate(&mut self) -> ManagementResult<Option<Vec<ComponentId>>> {
         let mut pending_operations = self.creates().await?;
         let mut component_graph = DependencyGraph::new();
         pending_operations.extend(self.updates().await?);
@@ -1029,6 +1029,9 @@ impl<'a> ManagementOperator<'a> {
             component_graph.id_depends_on(pending_parent.child_component_id, parent_id);
         }
 
+        let created_component_ids = (!self.created_components.is_empty())
+            .then_some(self.created_components.iter().copied().collect());
+
         self.send_component_ws_events(component_graph).await?;
 
         // Now, the rest of the pending ops can be executed, which need to have
@@ -1061,7 +1064,7 @@ impl<'a> ManagementOperator<'a> {
 
         self.actions().await?;
 
-        Ok(())
+        Ok(created_component_ids)
     }
 }
 
diff --git a/lib/dal/src/management/prototype.rs b/lib/dal/src/management/prototype.rs
index e5a28cbbeb..c80c3f1983 100644
--- a/lib/dal/src/management/prototype.rs
+++ b/lib/dal/src/management/prototype.rs
@@ -8,7 +8,7 @@ use std::{
 use serde::{Deserialize, Serialize};
 use si_events::FuncRunId;
 use thiserror::Error;
-use veritech_client::ManagementResultSuccess;
+use veritech_client::{ManagementFuncStatus, ManagementResultSuccess};
 
 use crate::{
     cached_module::{CachedModule, CachedModuleError},
@@ -24,7 +24,7 @@ use crate::{
     Component, ComponentError, ComponentId, DalContext, EdgeWeightKind,
     EdgeWeightKindDiscriminants, FuncId, HelperError, NodeWeightDiscriminants, Schema, SchemaError,
     SchemaId, SchemaVariant, SchemaVariantError, SchemaVariantId, TransactionsError,
-    WorkspaceSnapshotError, WsEvent, WsEventError, WsEventResult,
+    WorkspaceSnapshotError, WsEvent, WsEventError, WsEventResult, WsPayload,
 };
 
 use super::ManagementGeometry;
@@ -660,6 +660,16 @@ pub struct ManagementFuncExecutedPayload {
     func_run_id: FuncRunId,
 }
 
+#[derive(Clone, Deserialize, Serialize, Debug, PartialEq, Eq)]
+#[serde(rename_all = "camelCase")]
+pub struct ManagementOperationsCompletePayload {
+    request_ulid: Option<ulid::Ulid>,
+    func_name: String,
+    status: ManagementFuncStatus,
+    message: Option<String>,
+    created_component_ids: Option<Vec<ComponentId>>,
+}
+
 impl WsEvent {
     pub async fn management_func_executed(
         ctx: &DalContext,
@@ -669,7 +679,7 @@ impl WsEvent {
     ) -> WsEventResult<Self> {
         WsEvent::new(
             ctx,
-            crate::WsPayload::ManagementFuncExecuted(ManagementFuncExecutedPayload {
+            WsPayload::ManagementFuncExecuted(ManagementFuncExecutedPayload {
                 prototype_id,
                 manager_component_id,
                 func_run_id,
@@ -677,4 +687,25 @@ impl WsEvent {
         )
         .await
     }
+
+    pub async fn management_operations_complete(
+        ctx: &DalContext,
+        request_ulid: Option<ulid::Ulid>,
+        func_name: String,
+        message: Option<String>,
+        status: ManagementFuncStatus,
+        created_component_ids: Option<Vec<ComponentId>>,
+    ) -> WsEventResult<Self> {
+        WsEvent::new(
+            ctx,
+            WsPayload::ManagementOperationsComplete(ManagementOperationsCompletePayload {
+                request_ulid,
+                func_name,
+                status,
+                message,
+                created_component_ids,
+            }),
+        )
+        .await
+    }
 }
diff --git a/lib/dal/src/ws_event.rs b/lib/dal/src/ws_event.rs
index c12cb54321..ee760ff914 100644
--- a/lib/dal/src/ws_event.rs
+++ b/lib/dal/src/ws_event.rs
@@ -25,7 +25,9 @@ use crate::func::runner::FuncRunLogUpdatedPayload;
 use crate::func::{
     FuncWsEventCodeSaved, FuncWsEventFuncSummary, FuncWsEventGenerating, FuncWsEventPayload,
 };
-use crate::management::prototype::ManagementFuncExecutedPayload;
+use crate::management::prototype::{
+    ManagementFuncExecutedPayload, ManagementOperationsCompletePayload,
+};
 use crate::pkg::{
     ImportWorkspaceVotePayload, WorkspaceActorPayload, WorkspaceImportApprovalActorPayload,
 };
@@ -113,6 +115,7 @@ pub enum WsPayload {
     InferredEdgeRemove(InferredEdgeRemovePayload),
     InferredEdgeUpsert(InferredEdgeUpsertPayload),
     ManagementFuncExecuted(ManagementFuncExecutedPayload),
+    ManagementOperationsComplete(ManagementOperationsCompletePayload),
     ModuleImported(Vec<si_frontend_types::SchemaVariant>),
     Online(OnlinePayload),
     PromptUpdated(PromptUpdatedPayload),
diff --git a/lib/sdf-server/src/service/v2/management.rs b/lib/sdf-server/src/service/v2/management.rs
index 9aeed4730f..97f8cc26c4 100644
--- a/lib/sdf-server/src/service/v2/management.rs
+++ b/lib/sdf-server/src/service/v2/management.rs
@@ -7,7 +7,7 @@ use axum::{
     http::StatusCode,
     response::{IntoResponse, Response},
     routing::{get, post},
-    Router,
+    Json, Router,
 };
 use dal::{
     diagram::view::ViewId,
@@ -17,8 +17,8 @@ use dal::{
         ManagementError, ManagementFuncReturn, ManagementOperator,
     },
     schema::variant::authoring::VariantAuthoringError,
-    ChangeSet, ChangeSetError, ChangeSetId, ComponentId, FuncId, SchemaVariantError,
-    TransactionsError, WorkspacePk, WsEventError,
+    ChangeSet, ChangeSetError, ChangeSetId, ComponentId, Func, FuncError, FuncId,
+    SchemaVariantError, TransactionsError, WorkspacePk, WsEvent, WsEventError,
 };
 use serde::{Deserialize, Serialize};
 use si_layer_cache::LayerDbError;
@@ -37,13 +37,6 @@ mod generate_template;
 mod history;
 mod latest;
 
-#[derive(Debug, Deserialize, Serialize)]
-#[serde(rename_all = "camelCase")]
-pub struct RunPrototypeResponse {
-    status: ManagementFuncStatus,
-    message: Option<String>,
-}
-
 pub type ManagementApiResult<T> = Result<T, ManagementApiError>;
 
 #[remain::sorted]
@@ -51,6 +44,8 @@ pub type ManagementApiResult<T> = Result<T, ManagementApiError>;
 pub enum ManagementApiError {
     #[error("change set error: {0}")]
     ChangeSet(#[from] ChangeSetError),
+    #[error("func error: {0}")]
+    Func(#[from] FuncError),
     #[error("func api error: {0}")]
     FuncAPI(#[from] FuncAPIError),
     #[error("func authoring error: {0}")]
@@ -87,6 +82,19 @@ impl IntoResponse for ManagementApiError {
     }
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct RunPrototypeRequest {
+    request_ulid: Option<ulid::Ulid>,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct RunPrototypeResponse {
+    status: ManagementFuncStatus,
+    message: Option<String>,
+}
+
 pub async fn run_prototype(
     HandlerContext(builder): HandlerContext,
     AccessBuilder(access_builder): AccessBuilder,
@@ -100,6 +108,7 @@ pub async fn run_prototype(
         ComponentId,
         ViewId,
     )>,
+    Json(request): Json<RunPrototypeRequest>,
 ) -> ManagementApiResult<ForceChangeSetResponse<RunPrototypeResponse>> {
     let mut ctx = builder
         .build(access_builder.build(change_set_id.into()))
@@ -128,9 +137,10 @@ pub async fn run_prototype(
 
     if let Some(result) = execution_result.result.take() {
         let result: ManagementFuncReturn = result.try_into()?;
+        let mut created_component_ids = None;
         if result.status == ManagementFuncStatus::Ok {
             if let Some(operations) = result.operations {
-                ManagementOperator::new(
+                created_component_ids = ManagementOperator::new(
                     &ctx,
                     component_id,
                     operations,
@@ -143,6 +153,21 @@ pub async fn run_prototype(
             }
         }
 
+        let func_id = ManagementPrototype::func_id(&ctx, prototype_id).await?;
+        let func = Func::get_by_id_or_error(&ctx, func_id).await?;
+
+        WsEvent::management_operations_complete(
+            &ctx,
+            request.request_ulid,
+            func.name,
+            result.message.clone(),
+            result.status,
+            created_component_ids,
+        )
+        .await?
+        .publish_on_commit(&ctx)
+        .await?;
+
         ctx.commit().await?;
 
         return Ok(ForceChangeSetResponse::new(