diff --git a/CHANGELOG.md b/CHANGELOG.md
index aef1facb..c776b2c4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - feat: support navigating to discover in alerting popover([#316](https://github.com/opensearch-project/dashboards-assistant/pull/316))
 - fix: incorrect string escaping of vega schema([325](https://github.com/opensearch-project/dashboards-assistant/pull/325))
 - feat: register the AI actions to query controls in discover([#327](https://github.com/opensearch-project/dashboards-assistant/pull/327))
+- fix: t2viz ux improvements([#330](https://github.com/opensearch-project/dashboards-assistant/pull/330))
 
 ### 📈 Features/Enhancements
 
diff --git a/common/types/config.ts b/common/types/config.ts
index 74531362..2b92644f 100644
--- a/common/types/config.ts
+++ b/common/types/config.ts
@@ -26,7 +26,8 @@ export const configSchema = schema.object({
     enabled: schema.boolean({ defaultValue: false }),
   }),
   branding: schema.object({
-    label: schema.string({ defaultValue: '' }),
+    label: schema.maybe(schema.string()),
+    logo: schema.maybe(schema.string()),
   }),
 });
 
diff --git a/public/components/ui_action_context_menu.tsx b/public/components/ui_action_context_menu.tsx
index 0e5b9f14..e15fbed6 100644
--- a/public/components/ui_action_context_menu.tsx
+++ b/public/components/ui_action_context_menu.tsx
@@ -31,6 +31,7 @@ export const ActionContextMenu = (props: Props) => {
           trigger: AI_ASSISTANT_QUERY_EDITOR_TRIGGER as any,
         })),
         closeMenu: () => setOpen(false),
+        title: '',
       }),
     []
   );
diff --git a/public/components/visualization/source_selector.tsx b/public/components/visualization/source_selector.tsx
index 90438a4e..3cbb620d 100644
--- a/public/components/visualization/source_selector.tsx
+++ b/public/components/visualization/source_selector.tsx
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useCallback, useMemo, useState, useEffect } from 'react';
+import React, { useCallback, useMemo, useState, useEffect, useRef } from 'react';
 import { i18n } from '@osd/i18n';
 
 import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
@@ -38,6 +38,8 @@ export const SourceSelector = ({
   } = useOpenSearchDashboards<StartServices>();
   const [currentDataSources, setCurrentDataSources] = useState<DataSource[]>([]);
   const [dataSourceOptions, setDataSourceOptions] = useState<DataSourceGroup[]>([]);
+  const onChangeRef = useRef(onChange);
+  onChangeRef.current = onChange;
 
   const selectedSources = useMemo(() => {
     if (selectedSourceId) {
@@ -52,13 +54,20 @@ export const SourceSelector = ({
     return [];
   }, [selectedSourceId, dataSourceOptions]);
 
+  /**
+   * When initialized, select the first non-disabled option
+   */
   useEffect(() => {
     if (
       !selectedSourceId &&
       dataSourceOptions.length > 0 &&
       dataSourceOptions[0].options.length > 0
     ) {
-      onChange(dataSourceOptions[0].options[0]);
+      const options = dataSourceOptions[0].options;
+      const selectedOption = options.find((o) => !o.disabled);
+      if (selectedOption) {
+        onChangeRef.current(selectedOption);
+      }
     }
   }, [selectedSourceId, dataSourceOptions]);
 
@@ -81,7 +90,7 @@ export const SourceSelector = ({
 
   const onSetDataSourceOptions = useCallback(
     async (options: DataSourceGroup[]) => {
-      // Only support opensearch default data source
+      // Only support OpenSearch default data source
       const indexPatternOptions = options.find(
         (item) => item.groupType === DEFAULT_DATA_SOURCE_TYPE
       );
@@ -117,28 +126,30 @@ export const SourceSelector = ({
        * Check each data source to see if text to vega agents are configured or not
        * If not configured, disable the corresponding index pattern from the selection list
        */
-      Object.keys(dataSourceIdToIndexPatternIds).forEach(async (key) => {
-        const res = await assistantService.client.agentConfigExists(
-          [
-            TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
-            TEXT2VEGA_WITH_INSTRUCTIONS_AGENT_CONFIG_ID,
-            TEXT2PPL_AGENT_CONFIG_ID,
-          ],
-          {
-            dataSourceId: key !== 'DEFAULT' ? key : undefined,
-          }
-        );
-        if (!res.exists) {
-          dataSourceIdToIndexPatternIds[key].forEach((indexPatternId) => {
-            indexPatternOptions.options.forEach((option) => {
-              if (option.value === indexPatternId) {
-                option.disabled = true;
-              }
+      const updateIndexPatternPromises = Object.keys(dataSourceIdToIndexPatternIds).map(
+        async (key) => {
+          const res = await assistantService.client.agentConfigExists(
+            [
+              TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
+              TEXT2VEGA_WITH_INSTRUCTIONS_AGENT_CONFIG_ID,
+              TEXT2PPL_AGENT_CONFIG_ID,
+            ],
+            {
+              dataSourceId: key !== 'DEFAULT' ? key : undefined,
+            }
+          );
+          if (!res.exists) {
+            dataSourceIdToIndexPatternIds[key].forEach((indexPatternId) => {
+              indexPatternOptions.options.forEach((option) => {
+                if (option.value === indexPatternId) {
+                  option.disabled = true;
+                }
+              });
             });
-          });
+          }
         }
-      });
-
+      );
+      await Promise.allSettled(updateIndexPatternPromises);
       setDataSourceOptions([indexPatternOptions]);
     },
     [currentDataSources]
diff --git a/public/components/visualization/text2vega.ts b/public/components/visualization/text2vega.ts
index 44964971..aba0e388 100644
--- a/public/components/visualization/text2vega.ts
+++ b/public/components/visualization/text2vega.ts
@@ -63,6 +63,9 @@ export class Text2Vega {
                 )
                 // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 .toPromise<any>();
+              if (res.rawResponse.total === 0) {
+                throw new Error(`There is no result with the generated query: '${value.ppl}'.`);
+              }
               return { ...value, sample: res.rawResponse };
             }),
             // call llm to generate vega
diff --git a/public/components/visualization/text2viz.tsx b/public/components/visualization/text2viz.tsx
index a3b2ad78..99ee35e8 100644
--- a/public/components/visualization/text2viz.tsx
+++ b/public/components/visualization/text2viz.tsx
@@ -67,6 +67,7 @@ export const Text2Viz = () => {
       data,
       uiSettings,
       savedObjects,
+      config,
     },
   } = useOpenSearchDashboards<StartServices>();
 
@@ -106,7 +107,7 @@ export const Text2Viz = () => {
         if (result.error) {
           notifications.toasts.addError(result.error, {
             title: i18n.translate('dashboardAssistant.feature.text2viz.error', {
-              defaultMessage: 'Error while executing text to vega',
+              defaultMessage: 'Error while executing text to visualization',
             }),
           });
         } else {
@@ -198,7 +199,7 @@ export const Text2Viz = () => {
     });
 
     setSubmitting(false);
-  }, [selectedSource, input, status]);
+  }, [selectedSource, input, status, notifications.toasts]);
 
   /**
    * Display the save visualization dialog to persist the current generated visualization
@@ -342,16 +343,17 @@ export const Text2Viz = () => {
             onChange={(e) => setInput(e.target.value)}
             fullWidth
             compressed
-            prepend={<EuiIcon type={chatIcon} />}
+            prepend={<EuiIcon type={config.branding.logo || chatIcon} />}
             placeholder="Generate visualization with a natural language question."
             onKeyDown={(e) => e.key === 'Enter' && onSubmit()}
+            disabled={!selectedSource}
           />
         </EuiFlexItem>
         <EuiFlexItem grow={false}>
           <EuiButtonIcon
             aria-label="submit"
             onClick={onSubmit}
-            isDisabled={loading || input.trim().length === 0}
+            isDisabled={loading || input.trim().length === 0 || !selectedSource}
             display="base"
             size="s"
             iconType="returnKey"
diff --git a/public/plugin.tsx b/public/plugin.tsx
index f220df5d..fc18eb60 100644
--- a/public/plugin.tsx
+++ b/public/plugin.tsx
@@ -200,6 +200,7 @@ export class AssistantPlugin
               ...pluginsStart,
               ...coreStart,
               setHeaderActionMenu: params.setHeaderActionMenu,
+              config: this.config,
             });
 
             return () => {
diff --git a/public/types.ts b/public/types.ts
index f6c279ae..317f78c2 100644
--- a/public/types.ts
+++ b/public/types.ts
@@ -22,6 +22,7 @@ import { ExpressionsSetup, ExpressionsStart } from '../../../src/plugins/express
 import { SavedObjectsStart } from '../../../src/plugins/saved_objects/public';
 
 import { UsageCollectionSetup } from '../../../src/plugins/usage_collection/public';
+import { ConfigSchema } from '../common/types/config';
 
 export interface RenderProps {
   props: MessageContentProps;
@@ -95,6 +96,7 @@ export interface AssistantStart {
 export type StartServices = CoreStart &
   Omit<AssistantPluginStartDependencies, 'savedObjects'> & {
     setHeaderActionMenu: AppMountParameters['setHeaderActionMenu'];
+    config: ConfigSchema;
   };
 
 export interface UserAccount {