From b1bc05fb1d8e885b902ae56ebd39dbf1c3e429fb Mon Sep 17 00:00:00 2001
From: Nicolas Borges <nickdb@amazon.com>
Date: Tue, 7 Jul 2026 12:23:23 -0400
Subject: [PATCH] fix: fix infinite loop between source type and add another KB
 TUI screens

---
 .../screens/knowledge-base/AddKnowledgeBaseScreen.tsx | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/cli/tui/screens/knowledge-base/AddKnowledgeBaseScreen.tsx b/src/cli/tui/screens/knowledge-base/AddKnowledgeBaseScreen.tsx
index 14596c8f1..74317367c 100644
--- a/src/cli/tui/screens/knowledge-base/AddKnowledgeBaseScreen.tsx
+++ b/src/cli/tui/screens/knowledge-base/AddKnowledgeBaseScreen.tsx
@@ -212,10 +212,7 @@ export function AddKnowledgeBaseScreen({
       setPendingType(item.id as DataSourceTypeFlag);
       setStep('sources');
     },
-    // Esc from the type picker: if we already have at least one captured
-    // source, the only sensible return is the add-another decision (we can't
-    // un-capture earlier sources). Otherwise go back to description.
-    onExit: () => setStep(dataSources.length === 0 ? 'description' : 'add-another'),
+    onExit: () => setStep('description'),
   });
 
   // Surface the "Remove a captured source" option only when there's something
@@ -381,7 +378,7 @@ export function AddKnowledgeBaseScreen({
             prompt="Description (optional, press Enter to skip)"
             onSubmit={(value: string) => {
               setDescription(value);
-              setStep('data-source-type');
+              setStep(dataSources.length > 0 ? 'add-another' : 'data-source-type');
             }}
             onCancel={() => setStep('name')}
             allowEmpty
@@ -413,7 +410,7 @@ export function AddKnowledgeBaseScreen({
               setDataSources([...dataSources, { dataSourceType: pendingType, value }]);
               setStep('add-another');
             }}
-            onCancel={() => setStep(dataSources.length === 0 ? 'data-source-type' : 'add-another')}
+            onCancel={() => setStep('data-source-type')}
             schema={S3UriSchema}
           />
         )}
@@ -442,7 +439,7 @@ export function AddKnowledgeBaseScreen({
               ]);
               setStep('add-another');
             }}
-            onCancel={() => setStep(dataSources.length === 0 ? 'data-source-type' : 'add-another')}
+            onCancel={() => setStep('data-source-type')}
             schema={makeConnectorConfigSchema(pendingType)}
           />
         )}