From 269aa13bba321c6ba9e367ef8a54335ef83a5fc5 Mon Sep 17 00:00:00 2001
From: arc53-machine <232052973+arc53-machine@users.noreply.github.com>
Date: Mon, 28 Sep 2026 18:45:27 +0100
Subject: [PATCH] Address review: tool group order in the agent builder, saved
 S3/Reddit account

- Connection groups sit between the default and custom tools; a second
  sort that did not know connection labels moved them to the top.
- Upload picks the first saved account when connections load after the
  modal opened on S3 or Reddit, without overriding the user's choice.
---
 frontend/src/agents/NewAgent.test.tsx | 52 +++++++++++++++++++++++++--
 frontend/src/agents/NewAgent.tsx      | 16 ++++-----
 frontend/src/upload/Upload.test.tsx   | 42 ++++++++++++++++++++++
 frontend/src/upload/Upload.tsx        | 25 ++++++++-----
 4 files changed, 114 insertions(+), 21 deletions(-)

diff --git a/frontend/src/agents/NewAgent.test.tsx b/frontend/src/agents/NewAgent.test.tsx
index 5bd59f13..d6ffa3b8 100644
--- a/frontend/src/agents/NewAgent.test.tsx
+++ b/frontend/src/agents/NewAgent.test.tsx
@@ -27,6 +27,8 @@ const mocks = vi.hoisted(() => {
     dispatch: vi.fn(),
     getAgent: vi.fn(() => jsonResponse({})),
     createAgent: vi.fn(() => jsonResponse({ message: 'Name is taken' }, false)),
+    tools: null as unknown[] | null,
+    connections: [] as unknown[],
   };
 });
 const { jsonResponse } = mocks;
@@ -40,7 +42,7 @@ vi.mock('../api/services/userService', () => ({
   default: {
     getUserTools: () =>
       jsonResponse({
-        tools: [
+        tools: mocks.tools ?? [
           {
             id: 'tool-1',
             name: 'remote_device',
@@ -69,7 +71,7 @@ vi.mock('../api/services/devicesService', () => ({
 
 vi.mock('../api/services/connectorsService', () => ({
   default: {
-    listConnections: () => Promise.resolve({ connections: [] }),
+    listConnections: () => Promise.resolve({ connections: mocks.connections }),
   },
 }));
 
@@ -93,7 +95,9 @@ vi.mock('../components/MultiSelectPopover', () => ({
     <div data-testid="picker">
       {trigger}
       {items.map((item) => (
-        <div key={item.id}>{item.descriptionNode}</div>
+        <div key={item.id} data-group={item.group}>
+          {item.descriptionNode}
+        </div>
       ))}
     </div>
   ),
@@ -164,6 +168,8 @@ describe('NewAgent form', () => {
     mocks.dispatch.mockClear();
     mocks.getAgent.mockClear();
     mocks.createAgent.mockClear();
+    mocks.tools = null;
+    mocks.connections = [];
   });
 
   const render = async () => {
@@ -251,6 +257,46 @@ describe('NewAgent form', () => {
     expect(description.className).toContain('sm:col-start-2');
   });
 
+  it('lists tool groups as built-in, default, one per connection, then custom', async () => {
+    mocks.tools = [
+      { id: 'custom', name: 'api_tool', display_name: 'My API' },
+      {
+        id: 'linear',
+        name: 'mcp_tool',
+        display_name: 'Linear',
+        connection_id: 'c-lin',
+      },
+      { id: 'memory', name: 'memory', display_name: 'Memory', builtin: true },
+      {
+        id: 'notion',
+        name: 'mcp_tool',
+        display_name: 'Notion',
+        connection_id: 'c-not',
+      },
+      {
+        id: 'reader',
+        name: 'read_webpage',
+        display_name: 'Reader',
+        default: true,
+      },
+    ];
+    mocks.connections = [
+      { id: 'c-lin', name: 'Linear', account_label: 'a@x', icon: 'linear' },
+      { id: 'c-not', name: 'Notion', account_label: 'b@x', icon: 'notion' },
+    ];
+    await render();
+    const groups = Array.from(
+      container.querySelectorAll('[data-testid="picker"] [data-group]'),
+    ).map((item) => item.getAttribute('data-group'));
+    const order = groups.filter((g, i) => groups.indexOf(g) === i);
+    expect(order).toEqual([
+      'agents.form.toolsPopup.groupBuiltin',
+      'agents.form.toolsPopup.groupDefault',
+      'agents.form.toolsPopup.groupConnection',
+      'agents.form.toolsPopup.groupCustom',
+    ]);
+  });
+
   it('labels every picker with a floating label', async () => {
     await render();
     const labels = Array.from(
diff --git a/frontend/src/agents/NewAgent.tsx b/frontend/src/agents/NewAgent.tsx
index ba2719a0..92b68320 100644
--- a/frontend/src/agents/NewAgent.tsx
+++ b/frontend/src/agents/NewAgent.tsx
@@ -636,7 +636,12 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) {
         return t('agents.form.toolsPopup.groupCustom');
       };
       const tools: MultiSelectPopoverItem[] = [...visibleTools]
-        .sort((a, b) => rank(a) - rank(b))
+        .sort(
+          (a, b) =>
+            rank(a) - rank(b) ||
+            // Keeps each connection's tools together.
+            (rank(a) === 2 ? groupFor(a).localeCompare(groupFor(b)) : 0),
+        )
         .map((tool: UserToolType) => {
           const connection = connectionOf(tool);
           const base: MultiSelectPopoverItem = {
@@ -666,15 +671,6 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) {
           }
           return base;
         });
-      const groupOrder = [
-        t('agents.form.toolsPopup.groupBuiltin'),
-        t('agents.form.toolsPopup.groupDefault'),
-        t('agents.form.toolsPopup.groupCustom'),
-      ];
-      tools.sort(
-        (a, b) =>
-          groupOrder.indexOf(a.group || '') - groupOrder.indexOf(b.group || ''),
-      );
       setUserTools(tools);
       setRawUserTools(visibleTools);
     };
diff --git a/frontend/src/upload/Upload.test.tsx b/frontend/src/upload/Upload.test.tsx
index 457a8777..1373ef84 100644
--- a/frontend/src/upload/Upload.test.tsx
+++ b/frontend/src/upload/Upload.test.tsx
@@ -169,6 +169,48 @@ describe('Upload source-type tiles', () => {
     connectorsState.connections = [];
   });
 
+  it('picks the saved account once connections finish loading', async () => {
+    connectorsState.catalog = [
+      {
+        key: 's3',
+        icon: 's3',
+        sync_ingestor: 's3',
+        auth_kind: 'api_key',
+        available: true,
+        missing_settings: [],
+        credential_fields: [
+          { key: 'aws_access_key_id', label: 'Access key ID', secret: false },
+        ],
+      },
+    ];
+    const props = {
+      receivedFile: [],
+      setModalState: vi.fn(),
+      isOnboarding: false,
+      renderTab: null,
+      close: vi.fn(),
+      initialIngestor: 's3' as const,
+    };
+    await act(async () => root.render(<Upload {...props} />));
+    const triggers = () =>
+      Array.from(
+        document.body.querySelectorAll('[data-slot="select-trigger"]'),
+      ).map((el) => el.textContent);
+    expect(triggers()).not.toContain('modals.uploadDoc.newCredentials');
+    connectorsState.connections = [
+      {
+        id: 'k1',
+        connector_key: 's3',
+        status: 'connected',
+        account_label: '…WXYZ',
+      },
+    ];
+    await act(async () => root.render(<Upload {...props} />));
+    expect(triggers()).toContain('settings.connectors.detail.keyEnding');
+    connectorsState.catalog = [];
+    connectorsState.connections = [];
+  });
+
   it('leaves the disabled Train button on the default variant', async () => {
     await render();
     const crawler = tiles().find((tile) =>
diff --git a/frontend/src/upload/Upload.tsx b/frontend/src/upload/Upload.tsx
index 418004c0..22c5e4e5 100644
--- a/frontend/src/upload/Upload.tsx
+++ b/frontend/src/upload/Upload.tsx
@@ -1,7 +1,7 @@
 import { ChevronLeft, CircleAlert, FileText } from 'lucide-react';
 import { envVar } from '@/env';
 import { cn } from '@/lib/utils';
-import { useCallback, useEffect, useState } from 'react';
+import { useCallback, useEffect, useRef, useState } from 'react';
 import { nanoid } from '@reduxjs/toolkit';
 import type { FileRejection } from 'react-dropzone';
 import { useTranslation } from 'react-i18next';
@@ -179,7 +179,10 @@ function Upload({
           <UiFormField label={t('filePicker.account')}>
             <Select
               value={connectionId ?? 'new'}
-              onValueChange={(value) => setConnectionId(value)}
+              onValueChange={(value) => {
+                accountPicked.current = true;
+                setConnectionId(value);
+              }}
             >
               <SelectTrigger className="w-full" size="field" shape="pill">
                 <SelectValue />
@@ -538,14 +541,20 @@ function Upload({
   const usingSavedKeys =
     keyAccounts.length > 0 && !!connectionId && connectionId !== 'new';
 
+  // Default to the first saved account, also when the connections arrive
+  // after the modal opened on S3 or Reddit, but never over a choice the
+  // user made in the account picker.
+  const accountPicked = useRef(false);
+  const firstKeyAccount = keyAccounts[0]?.id;
   useEffect(() => {
-    if (selectedConnector?.auth_kind === 'api_key') {
-      setConnectionId(keyAccounts[0]?.id ?? 'new');
-    } else {
-      setConnectionId(null);
-    }
-    // Reset only when the source type changes.
+    accountPicked.current = false;
+    setConnectionId(null);
   }, [ingestor.type]);
+  useEffect(() => {
+    if (selectedConnector?.auth_kind !== 'api_key' || accountPicked.current)
+      return;
+    setConnectionId(firstKeyAccount ?? 'new');
+  }, [ingestor.type, selectedConnector?.auth_kind, firstKeyAccount]);
 
   const ingestorOptions: IngestorOption[] = IngestorFormSchemas.map(
     (schema) => ({