From 97a58bf5db4e82e59db4a15760a5849f04ce4df4 Mon Sep 17 00:00:00 2001
From: Nicolas Chaulet <nicolas.chaulet@elastic.co>
Date: Tue, 9 Jan 2024 15:05:01 -0500
Subject: [PATCH] [Fleet] Fix package edit redirect (#174531)

(cherry picked from commit 66aa37444a16c5d40523b9cae3267983dd20a9db)
---
 .../package_policies/package_policies_table.tsx |  1 +
 .../sections/epm/screens/policy/index.tsx       | 17 +++++++++++++++--
 .../components/package_policy_actions_menu.tsx  |  6 ++++--
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/details_page/components/package_policies/package_policies_table.tsx b/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/details_page/components/package_policies/package_policies_table.tsx
index 6f835651e6dbf..934fee64b4043 100644
--- a/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/details_page/components/package_policies/package_policies_table.tsx
+++ b/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/details_page/components/package_policies/package_policies_table.tsx
@@ -238,6 +238,7 @@ export const PackagePoliciesTable: React.FunctionComponent<Props> = ({
               return canWriteIntegrationPolicies ? (
                 <PackagePolicyActionsMenu
                   agentPolicy={agentPolicy}
+                  from={'fleet-policy-list'}
                   packagePolicy={packagePolicy}
                   upgradePackagePolicyHref={`${getHref('upgrade_package_policy', {
                     policyId: agentPolicy.id,
diff --git a/x-pack/plugins/fleet/public/applications/integrations/sections/epm/screens/policy/index.tsx b/x-pack/plugins/fleet/public/applications/integrations/sections/epm/screens/policy/index.tsx
index 8eeb8622a56a7..7bafedc453e17 100644
--- a/x-pack/plugins/fleet/public/applications/integrations/sections/epm/screens/policy/index.tsx
+++ b/x-pack/plugins/fleet/public/applications/integrations/sections/epm/screens/policy/index.tsx
@@ -6,10 +6,11 @@
  */
 
 import React, { memo } from 'react';
-import { useRouteMatch } from 'react-router-dom';
+import { useLocation, useRouteMatch } from 'react-router-dom';
 
 // TODO: Needs to be moved
 import { EditPackagePolicyForm } from '../../../../../fleet/sections/agent_policy/edit_package_policy_page';
+import type { EditPackagePolicyFrom } from '../../../../../fleet/sections/agent_policy/create_package_policy_page/types';
 import { useGetOnePackagePolicyQuery, useUIExtension } from '../../../../hooks';
 
 export const Policy = memo(() => {
@@ -17,6 +18,7 @@ export const Policy = memo(() => {
     params: { packagePolicyId },
   } = useRouteMatch<{ packagePolicyId: string }>();
 
+  const { search } = useLocation();
   const { data: packagePolicyData } = useGetOnePackagePolicyQuery(packagePolicyId);
 
   const extensionView = useUIExtension(
@@ -24,10 +26,21 @@ export const Policy = memo(() => {
     'package-policy-edit'
   );
 
+  const qs = new URLSearchParams(search);
+  const fromQs = qs.get('from');
+
+  let from: EditPackagePolicyFrom | undefined;
+
+  if (fromQs && fromQs === 'fleet-policy-list') {
+    from = 'edit';
+  } else {
+    from = 'package-edit';
+  }
+
   return (
     <EditPackagePolicyForm
       packagePolicyId={packagePolicyId}
-      from="package-edit"
+      from={from}
       forceUpgrade={extensionView?.useLatestPackageVersion}
     />
   );
diff --git a/x-pack/plugins/fleet/public/components/package_policy_actions_menu.tsx b/x-pack/plugins/fleet/public/components/package_policy_actions_menu.tsx
index 24525ce709efd..e2f581d38d1db 100644
--- a/x-pack/plugins/fleet/public/components/package_policy_actions_menu.tsx
+++ b/x-pack/plugins/fleet/public/components/package_policy_actions_menu.tsx
@@ -24,12 +24,14 @@ export const PackagePolicyActionsMenu: React.FunctionComponent<{
   showAddAgent?: boolean;
   defaultIsOpen?: boolean;
   upgradePackagePolicyHref?: string;
+  from?: 'fleet-policy-list' | undefined;
 }> = ({
   agentPolicy,
   packagePolicy,
   showAddAgent,
   upgradePackagePolicyHref,
   defaultIsOpen = false,
+  from,
 }) => {
   const [isEnrollmentFlyoutOpen, setIsEnrollmentFlyoutOpen] = useState(false);
   const { getHref } = useLink();
@@ -80,9 +82,9 @@ export const PackagePolicyActionsMenu: React.FunctionComponent<{
       data-test-subj="PackagePolicyActionsEditItem"
       disabled={!canWriteIntegrationPolicies || !agentPolicy}
       icon="pencil"
-      href={getHref('integration_policy_edit', {
+      href={`${getHref('integration_policy_edit', {
         packagePolicyId: packagePolicy.id,
-      })}
+      })}${from ? `?from=${from}` : ''}`}
       key="packagePolicyEdit"
     >
       <FormattedMessage