[AI Chat] UI - Better Tool/Task status (#32545)

* [AI Chat] UI - Better Tool/Task status - tool events only show active when they are executing - "Thinking" is shown when tool results are being sent and the response generated by the AI - "Paused" / "Stopped" is shown for the active task in those states - Browser test for these labels
2025-11-27 11:28:02 +01:00
parent b7f4a5afa0
commit cb15b48e1c
15 changed files with 372 additions and 34 deletions
@@ -18,6 +18,7 @@
 #include "base/test/scoped_feature_list.h"
 #include "base/test/test_future.h"
 #include "brave/browser/ai_chat/ai_chat_service_factory.h"
+#include "brave/browser/ui/webui/ai_chat/ai_chat_untrusted_conversation_ui.h"
 #include "brave/components/ai_chat/core/browser/ai_chat_service.h"
 #include "brave/components/ai_chat/core/browser/conversation_handler.h"
 #include "brave/components/ai_chat/core/browser/engine/mock_engine_consumer.h"
@@ -35,6 +36,7 @@
 #include "chrome/browser/ui/browser_navigator_params.h"
 #include "chrome/test/base/in_process_browser_test.h"
 #include "chrome/test/base/ui_test_utils.h"
+#include "content/public/browser/render_frame_host.h"
 #include "content/public/browser/web_contents.h"
 #include "content/public/test/browser_test.h"
 #include "content/public/test/browser_test_utils.h"
@@ -160,12 +162,43 @@ class AIChatConversationTaskBrowserTest : public InProcessBrowserTest {
    conversation_rfh_ = ui_test_utils::NavigateToURLWithDisposition(
        agent_browser_window_, url, WindowOpenDisposition::NEW_FOREGROUND_TAB,
        ui_test_utils::BROWSER_TEST_WAIT_FOR_LOAD_STOP);
+    // Wait for child frame to exist
+    VerifyElementState("conversation-entries-iframe");
+    EXPECT_TRUE(base::test::RunUntil(
+        [&] { return GetConversationEntriesFrame() != nullptr; }));
+  }
+
+  content::RenderFrameHost* GetConversationEntriesFrame() {
+    content::RenderFrameHost* result = nullptr;
+    conversation_rfh_->ForEachRenderFrameHost(
+        [&](content::RenderFrameHost* frame) {
+          if (frame->GetWebUI() &&
+              frame->GetWebUI()
+                  ->GetController()
+                  ->GetAs<AIChatUntrustedConversationUI>()) {
+            result = frame;
+            return;
+          }
+        });
+    return result;
+  }
+
+  bool VerifyConversationFrameElementState(
+      const std::string& test_id,
+      bool expect_exist = true,
+      base::Location location = base::Location::Current()) {
+    return VerifyElementState(test_id, expect_exist,
+                              GetConversationEntriesFrame(), location);
  }

  // Helper to check if an element with a specific data-testid exists
  bool VerifyElementState(const std::string& test_id,
                          bool expect_exist = true,
+                          content::RenderFrameHost* frame = nullptr,
                          base::Location location = base::Location::Current()) {
+    if (!frame) {
+      frame = conversation_rfh_;
+    }
    SCOPED_TRACE(testing::Message()
                 << "VerifyElementState: '" << test_id << "' called from "
                 << location.file_name() << ":" << location.line_number());
@@ -198,7 +231,7 @@ class AIChatConversationTaskBrowserTest : public InProcessBrowserTest {
    )";

    auto result = content::EvalJs(
-        conversation_rfh_,
+        frame,
        content::JsReplace(kWaitForAIChatRenderScript, test_id, !expect_exist));
    return result.ExtractBool();
  }
@@ -245,25 +278,30 @@ class AIChatConversationTaskBrowserTest : public InProcessBrowserTest {
  base::OnceClosure SetupMockGenerateAssistantResponse(
      EngineConsumer::GenerationDataCallback* out_data_callback,
      EngineConsumer::GenerationCompletedCallback* out_completed_callback,
+      testing::Sequence* sequence = nullptr,
      base::Location location = base::Location::Current()) {
    SCOPED_TRACE(testing::Message() << location.ToString());
    auto run_loop = std::make_unique<base::RunLoop>();
    auto on_generate_called = run_loop->QuitClosure();
-    EXPECT_CALL(*mock_engine_,
-                GenerateAssistantResponse(_, _, _, _, _, _, _, _, _))
-        .Description(base::StrCat(
-            {"GenerateAssistantResponse mocked from ", location.ToString()}))
-        .WillOnce([out_data_callback, out_completed_callback,
-                   on_called = std::move(on_generate_called)](
-                      PageContentsMap page_contents,
-                      const EngineConsumer::ConversationHistory& history,
-                      const std::string& selected_language, bool is_temporary,
-                      const std::vector<base::WeakPtr<Tool>>& provided_tools,
-                      std::optional<std::string_view> preferred_tool_name,
-                      mojom::ConversationCapability capability,
-                      EngineConsumer::GenerationDataCallback data_cb,
-                      EngineConsumer::GenerationCompletedCallback
-                          complete_cb) mutable {
+    auto& expect =
+        EXPECT_CALL(*mock_engine_,
+                    GenerateAssistantResponse(_, _, _, _, _, _, _, _, _))
+            .Description(base::StrCat({"GenerateAssistantResponse mocked from ",
+                                       location.ToString()}));
+    if (sequence) {
+      expect.InSequence(*sequence);
+    }
+    expect.WillOnce(
+        [out_data_callback, out_completed_callback,
+         on_called = std::move(on_generate_called)](
+            PageContentsMap page_contents,
+            const EngineConsumer::ConversationHistory& history,
+            const std::string& selected_language, bool is_temporary,
+            const std::vector<base::WeakPtr<Tool>>& provided_tools,
+            std::optional<std::string_view> preferred_tool_name,
+            mojom::ConversationCapability capability,
+            EngineConsumer::GenerationDataCallback data_cb,
+            EngineConsumer::GenerationCompletedCallback complete_cb) mutable {
          *out_data_callback = std::move(data_cb);
          *out_completed_callback = std::move(complete_cb);
          std::move(on_called).Run();
@@ -283,6 +321,12 @@ class AIChatConversationTaskBrowserTest : public InProcessBrowserTest {
                                    nullptr);
  }

+  mojom::ToolUseEventPtr CreateToolUseEvent(const std::string& tool_name,
+                                            const std::string& tool_id) {
+    return mojom::ToolUseEvent::New(tool_name, tool_id, "{}", std::nullopt,
+                                    nullptr);
+  }
+
  raw_ptr<Profile> agent_profile_ = nullptr;
  raw_ptr<content::RenderFrameHost> conversation_rfh_ = nullptr;
  raw_ptr<Browser> agent_browser_window_ = nullptr;
@@ -557,6 +601,166 @@ IN_PROC_BROWSER_TEST_F(AIChatConversationTaskBrowserTest, TaskStopAction) {
  }
 }

+IN_PROC_BROWSER_TEST_F(AIChatConversationTaskBrowserTest, TaskUI) {
+  // A task UI shows when there are 2 tool segments of a tool loop, i.e. the AI
+  // responds to a tool use result with another tool use request.
+  CreateConversationWithMockEngine();
+  std::string uuid = conversation_handler_->get_conversation_uuid();
+
+  NavigateToConversationUI(uuid);
+
+  // Inject our own Tool so that we can handle the tool execution and pause
+  auto* mock_tool = static_cast<NiceMock<MockTool>*>(
+      content_agent_tool_provider_->AddToolForTesting(
+          std::make_unique<NiceMock<MockTool>>("mock_tool", "Mock tool")));
+
+  testing::Sequence tool_call_seq;
+  base::OnceClosure tool_execute;
+  // Submit first message
+  {
+    EngineConsumer::GenerationDataCallback data_callback;
+    EngineConsumer::GenerationCompletedCallback completed_callback;
+    auto wait_for_generate = SetupMockGenerateAssistantResponse(
+        &data_callback, &completed_callback, &tool_call_seq);
+    conversation_handler_->SubmitHumanConversationEntry(
+        "Navigate to example.com", std::nullopt);
+    std::move(wait_for_generate).Run();
+    // Send first message response
+    // Simulate tool use event
+    data_callback.Run(EngineConsumer::GenerationResultData(
+        mojom::ConversationEntryEvent::NewToolUseEvent(
+            CreateToolUseEvent("mock_tool", "tool_id_1")),
+        std::nullopt));
+
+    EXPECT_CALL(*mock_tool, UseTool)
+        .InSequence(tool_call_seq)
+        .WillOnce(testing::WithArg<1>([&](Tool::UseToolCallback callback) {
+          // Wait unt the next round is setup to call the callback
+          tool_execute = base::BindOnce(
+              [](Tool::UseToolCallback callback) {
+                std::move(callback).Run(
+                    CreateContentBlocksForText("1st tool result"));
+              },
+              std::move(callback));
+        }));
+
+    // Complete first message response
+    std::move(completed_callback)
+        .Run(base::ok(
+            EngineConsumer::GenerationResultData(nullptr, std::nullopt)));
+  }
+
+  // No task UI should be shown with only one tool segment in the loop
+  EXPECT_FALSE(VerifyConversationFrameElementState("assistant-task", false));
+  // Handle the tool execution response with another tool use request.
+  {
+    EngineConsumer::GenerationDataCallback data_callback;
+    EngineConsumer::GenerationCompletedCallback completed_callback;
+    auto wait_for_generate = SetupMockGenerateAssistantResponse(
+        &data_callback, &completed_callback, &tool_call_seq);
+    std::move(tool_execute).Run();
+    std::move(wait_for_generate).Run();
+
+    data_callback.Run(EngineConsumer::GenerationResultData(
+        mojom::ConversationEntryEvent::NewCompletionEvent(
+            mojom::CompletionEvent::New("Hmm, I want a different thing")),
+        std::nullopt));
+    data_callback.Run(EngineConsumer::GenerationResultData(
+        mojom::ConversationEntryEvent::NewToolUseEvent(
+            CreateToolUseEvent("mock_tool", "tool_id_2")),
+        std::nullopt));
+
+    // When the tool is being executed, we can verify the UI state
+    EXPECT_CALL(*mock_tool, UseTool)
+        .InSequence(tool_call_seq)
+        .WillOnce(testing::WithArg<1>([&](Tool::UseToolCallback callback) {
+          EXPECT_TRUE(VerifyConversationFrameElementState("assistant-task"));
+          EXPECT_FALSE(VerifyConversationFrameElementState(
+              "tool-event-thinking", false));
+          tool_execute = base::BindOnce(
+              [](Tool::UseToolCallback callback) {
+                std::move(callback).Run(
+                    CreateContentBlocksForText("2nd tool result"));
+              },
+              std::move(callback));
+        }));
+
+    // Complete successful response
+    std::move(completed_callback)
+        .Run(base::ok(
+            EngineConsumer::GenerationResultData(nullptr, std::nullopt)));
+  }
+
+  // Handle the second tool execution response with pausing and verify UI label.
+  {
+    EngineConsumer::GenerationDataCallback data_callback;
+    EngineConsumer::GenerationCompletedCallback completed_callback;
+    auto wait_for_generate = SetupMockGenerateAssistantResponse(
+        &data_callback, &completed_callback, &tool_call_seq);
+
+    // Finish executing the tool
+    std::move(tool_execute).Run();
+
+    std::move(wait_for_generate).Run();
+
+    // Now we should be thinking
+    EXPECT_TRUE(VerifyConversationFrameElementState("tool-event-thinking"));
+
+    // Shouldn't call the tool again because we are pausing.
+    EXPECT_CALL(*mock_tool, UseTool).Times(0).InSequence(tool_call_seq);
+
+    // Pause the task
+    EXPECT_TRUE(ClickElement("pause-task-button"));
+
+    data_callback.Run(EngineConsumer::GenerationResultData(
+        mojom::ConversationEntryEvent::NewCompletionEvent(
+            mojom::CompletionEvent::New("Hmm, I want a different thing")),
+        std::nullopt));
+    data_callback.Run(EngineConsumer::GenerationResultData(
+        mojom::ConversationEntryEvent::NewToolUseEvent(
+            CreateToolUseEvent("mock_tool", "tool_id_3")),
+        std::nullopt));
+    // Complete successful response
+    std::move(completed_callback)
+        .Run(base::ok(
+            EngineConsumer::GenerationResultData(nullptr, std::nullopt)));
+  }
+
+  // The task should have a "paused" label
+  EXPECT_TRUE(
+      VerifyConversationFrameElementState("assistant-task-paused-label"));
+
+  // When we submit a new message, the task is no longer active. It should still
+  // exist but should not have its "paused" label.
+  {
+    EngineConsumer::GenerationDataCallback data_callback;
+    EngineConsumer::GenerationCompletedCallback completed_callback;
+    auto wait_for_generate = SetupMockGenerateAssistantResponse(
+        &data_callback, &completed_callback, &tool_call_seq);
+    conversation_handler_->SubmitHumanConversationEntry(
+        "Actually do something different", std::nullopt);
+
+    EXPECT_TRUE(base::test::RunUntil([this]() {
+      return GetConversationState()->tool_use_task_state ==
+             mojom::TaskState::kNone;
+    }));
+
+    std::move(wait_for_generate).Run();
+    // Simple response
+    data_callback.Run(EngineConsumer::GenerationResultData(
+        mojom::ConversationEntryEvent::NewCompletionEvent(
+            mojom::CompletionEvent::New("ok")),
+        std::nullopt));
+    // Complete first message response
+    std::move(completed_callback)
+        .Run(base::ok(
+            EngineConsumer::GenerationResultData(nullptr, std::nullopt)));
+  }
+  EXPECT_TRUE(VerifyConversationFrameElementState("assistant-task"));
+  EXPECT_FALSE(VerifyConversationFrameElementState(
+      "assistant-task-paused-label", false));
+}
+
 #endif  // BUILDFLAG(ENABLE_BRAVE_AI_CHAT_AGENT_PROFILE)

 }  // namespace ai_chat
@@ -53,6 +53,11 @@ class ContentAgentToolProvider : public ToolProvider,
    return task_tab_handle_;
  }

+  Tool* AddToolForTesting(std::unique_ptr<Tool> tool) {
+    tools_.push_back(std::move(tool));
+    return tools_.back().get();
+  }
+
 private:
  friend class BrowserToolsTest;
  friend class ContentAgentToolProviderTest;
@@ -1925,8 +1925,9 @@ ConversationHandler::GetStateForConversationEntries() {
  mojom::ConversationEntriesStatePtr entries_state =
      mojom::ConversationEntriesState::New();

-  entries_state->is_generating =
-      IsRequestInProgress() || is_tool_use_in_progress_;
+  entries_state->is_generating = IsRequestInProgress();
+  entries_state->is_tool_executing = is_tool_use_in_progress_;
+  entries_state->tool_use_task_state = tool_use_task_state_;
  entries_state->is_leo_model = is_leo_model;
  entries_state->all_models = std::move(models_copy);
  entries_state->current_model_key = model.key;
@@ -1997,8 +1998,7 @@ void ConversationHandler::OnSuggestedQuestionsChanged() {
 void ConversationHandler::OnAPIRequestInProgressChanged() {
  OnStateForConversationEntriesChanged();
  for (auto& client : conversation_ui_handlers_) {
-    client->OnAPIRequestInProgress(is_request_in_progress_ ||
-                                   is_tool_use_in_progress_);
+    client->OnAPIRequestInProgress(is_request_in_progress_);
  }
  for (auto& observer : observers_) {
    observer.OnRequestInProgressChanged(this, is_request_in_progress_);
@@ -2006,6 +2006,7 @@ void ConversationHandler::OnAPIRequestInProgressChanged() {
 }

 void ConversationHandler::OnToolUseTaskStateChanged() {
+  OnStateForConversationEntriesChanged();
  for (auto& client : conversation_ui_handlers_) {
    client->OnTaskStateChanged(tool_use_task_state_);
  }
@@ -365,6 +365,10 @@ struct Model {
 struct ConversationEntriesState {
  // Whether an answer generation is in progress
  bool is_generating;
+  // Whether the most recent tool use request is being executed
+  bool is_tool_executing;
+  // State of any active tool loop task
+  TaskState tool_use_task_state;
  // Whether the current model is a built-in Leo model
  bool is_leo_model;
  array<Model> all_models;
@@ -162,6 +162,7 @@ function ConversationEntries(props: ConversationEntriesProps) {
        + conversationContext.conversationUuid
      }
      ref={iframeRef}
+      data-testid='conversation-entries-iframe'
      onLoad={() => setHasLoaded(true)}
    />
  )
@@ -992,6 +992,7 @@ type CustomArgs = {
  totalTokens: number
  trimmedTokens: number
  isGenerating: boolean
+  isToolExecuting: boolean
  attachmentsDialog: 'tabs' | 'bookmarks' | 'history' | null
  isNewConversation: boolean
  generatedUrlToBeOpened: Url | undefined
@@ -1041,6 +1042,7 @@ const args: CustomArgs = {
  totalTokens: 0,
  trimmedTokens: 0,
  isGenerating: false,
+  isToolExecuting: false,
  attachmentsDialog: null,
  isNewConversation: false,
  generatedUrlToBeOpened: undefined,
@@ -1395,6 +1397,8 @@ function StoryContext(
    conversationHistory: conversationContext.conversationHistory,
    conversationCapability: Mojom.ConversationCapability.CONTENT_AGENT,
    isGenerating: conversationContext.isGenerating,
+    isToolExecuting: args.isToolExecuting,
+    toolUseTaskState: conversationContext.toolUseTaskState,
    isLeoModel: conversationContext.isCurrentModelLeo,
    contentUsedPercentage: options.args.shouldShowLongPageWarning ? 48 : 100,
    visualContentUsedPercentage: options.args.shouldShowLongVisualContentWarning
@@ -1546,6 +1550,7 @@ export const _ToolUse = {
            <ToolEvent
              key={event.toolUseEvent!.id}
              toolUseEvent={event.toolUseEvent!}
+              isExecuting={args.isToolExecuting}
              isEntryActive
            ></ToolEvent>
          ))}
@@ -85,6 +85,7 @@ function AssistantEvent(
  },
 ) {
  const { allowedLinks, event, isEntryInProgress, isLeoModel } = props
+  const context = useUntrustedConversationContext()

  if (event.completionEvent) {
    const numberedLinks =
@@ -139,6 +140,7 @@ function AssistantEvent(
      <ToolEvent
        toolUseEvent={props.event.toolUseEvent}
        isEntryActive={props.isEntryInteractivityAllowed}
+        isExecuting={context.isToolExecuting}
      />
    )
  }
@@ -29,7 +29,7 @@
    }
  }

-  &.isActive .toolLabel {
+  &.isExecuting .toolLabel {
    background: linear-gradient(
      90deg,
      #434fcf 0%,
@@ -5,6 +5,7 @@

 import * as React from 'react'
 import classnames from '$web-common/classnames'
+import { getLocale } from '$web-common/locale'
 import * as Mojom from '../../../common/mojom'
 import ToolPermissionChallenge from '../tool_permission_challenge/tool_permission_challenge'
 import { getToolLabel } from './get_tool_label'
@@ -20,6 +21,12 @@ interface Props {
   * what to do with partial input.
   */
  isEntryActive: boolean
+
+  /**
+   * Whether the tool use request is currently being executed. This is usually
+   * a short-lived state.
+   */
+  isExecuting: boolean
 }

 /**
@@ -158,11 +165,16 @@ export default function ToolEvent(props: Props) {

        const isExpandable = content.toolLabel && content.expandedContent

+        // Show as executing if this is an active entry, we are currently
+        // executing *a* tool, and this tool has no output yet.
+        const isExecuting =
+          props.isEntryActive && props.isExecuting && !props.toolUseEvent.output
+
        return (
          <div
            className={classnames(
              styles.toolUse,
-              props.isEntryActive && styles.isActive,
+              isExecuting && styles.isExecuting,
            )}
          >
            {isExpandable && (
@@ -205,3 +217,16 @@ function ToolContentExpandable(props: Props & ToolUseContent) {
    </>
  )
 }
+
+export function ToolEventThinking() {
+  return (
+    <div
+      className={classnames(styles.toolUse, styles.isExecuting)}
+      data-testid='tool-event-thinking'
+    >
+      <div className={styles.toolLabel}>
+        {getLocale(S.CHAT_UI_TOOL_LABEL_THINKING)}
+      </div>
+    </div>
+  )
+}
@@ -38,6 +38,10 @@
  }
 }

+.taskStateLabel {
+  margin-top: var(--leo-spacing-m);
+}
+
 .taskStep {
  --task-icon-size: 22px;
  --half-task-icon-size: calc(var(--task-icon-size) / 2);
@@ -10,17 +10,22 @@ import * as Mojom from '../../../common/mojom'
 import { taskConversationEntries } from '../../../page/stories/story_utils/history'
 import MockContext from '../../mock_untrusted_conversation_context'
 import AssistantTask from './assistant_task'
+import { getKeysForMojomEnum } from '$web-common/mojomUtils'

 type CustomArgs = {
  isActiveTask: boolean
  isGenerating: boolean
+  isToolExecuting: boolean
  hasContentThumbnail: boolean
+  toolUseTaskState: keyof typeof Mojom.TaskState
 }

 const args: CustomArgs = {
  isActiveTask: false,
  isGenerating: false,
+  isToolExecuting: false,
  hasContentThumbnail: false,
+  toolUseTaskState: 'kNone' satisfies keyof typeof Mojom.TaskState,
 }

 export const _AssistantTask = {
@@ -28,6 +33,9 @@ export const _AssistantTask = {
    return (
      <MockContext
        contentTaskTabId={args.hasContentThumbnail ? 1 : undefined}
+        isGenerating={args.isGenerating}
+        isToolExecuting={args.isToolExecuting}
+        toolUseTaskState={Mojom.TaskState[args.toolUseTaskState]}
        uiObserver={
          {
            thumbnailUpdated: {
@@ -47,7 +55,6 @@ export const _AssistantTask = {
        <AssistantTask
          assistantEntries={taskConversationEntries}
          isActiveTask={args.isActiveTask}
-          isGenerating={args.isGenerating}
          isLeoModel={true}
        />
      </MockContext>
@@ -58,6 +65,12 @@ export const _AssistantTask = {
 export default {
  title: 'Chat/Chat',
  component: AssistantTask,
-  argTypes: InferControlsFromArgs(args),
+  argTypes: {
+    ...InferControlsFromArgs(args),
+    toolUseTaskState: {
+      options: getKeysForMojomEnum(Mojom.TaskState),
+      control: { type: 'select' },
+    },
+  },
  args,
 } as Meta<typeof AssistantTask>
@@ -5,6 +5,7 @@

 import * as React from 'react'
 import Icon from '@brave/leo/react/icon'
+import Label from '@brave/leo/react/label'
 import ProgressRing from '@brave/leo/react/progressRing'
 import Tabs from '@brave/leo/react/tabs'
 import TabItem from '@brave/leo/react/tabItem'
@@ -13,7 +14,7 @@ import { getLocale } from '$web-common/locale'
 import * as Mojom from '../../../common/mojom'
 import { useUntrustedConversationContext } from '../../untrusted_conversation_context'
 import AssistantResponse from '../assistant_response'
-import ToolEvent from '../assistant_response/tool_event'
+import ToolEvent, { ToolEventThinking } from '../assistant_response/tool_event'
 import styles from './assistant_task.module.scss'
 import useExtractTaskData, { TaskData } from './use_extract_task_data'

@@ -25,12 +26,25 @@ interface Props {
  // the conversation). Informs whether interactivity is allowed.
  isActiveTask: boolean

+  // Passes to AssistantResponse
+  isLeoModel: boolean
+}
+
+interface TabProps {
  // Whether the task is actively still being generated. Informs whether the UI
  // displays as if it's still in progress.
  isGenerating: boolean

-  // Passes to AssistantResponse
-  isLeoModel: boolean
+  // Whether the most recent tool use request is being executed
+  isToolExecuting: boolean
+
+  // Whether the task is currently generating a response
+  isThinking: boolean
+
+  // What the state of the active task is
+  toolUseTaskState: Mojom.TaskState
+
+  taskData: TaskData
 }

 /**
@@ -84,8 +98,24 @@ export default function AssistantTask(props: Props) {

  const taskData = useExtractTaskData(props.assistantEntries)

+  const isThinking =
+    props.isActiveTask
+    && !conversationContext.isToolExecuting
+    && conversationContext.toolUseTaskState === Mojom.TaskState.kRunning
+
+  const tabProps: TabProps = {
+    isGenerating: conversationContext.isGenerating,
+    isToolExecuting: conversationContext.isToolExecuting,
+    toolUseTaskState: conversationContext.toolUseTaskState,
+    isThinking: isThinking,
+    taskData: taskData,
+  }
+
  return (
-    <div className={styles.task}>
+    <div
+      className={styles.task}
+      data-testid='assistant-task'
+    >
      <Tabs
        onChange={() => setShowSteps(!showSteps)}
        value={showSteps ? 'steps' : 'progress'}
@@ -101,16 +131,44 @@ export default function AssistantTask(props: Props) {
          {showSteps && (
            <Steps
              {...props}
-              taskData={taskData}
+              {...tabProps}
            />
          )}

          {!showSteps && (
            <Progress
              {...props}
-              taskData={taskData}
+              {...tabProps}
            />
          )}
+
+          {props.isActiveTask
+            && conversationContext.toolUseTaskState
+              === Mojom.TaskState.kPaused && (
+              <Label
+                color='neutral'
+                mode='loud'
+                className={styles.taskStateLabel}
+              >
+                <span data-testid='assistant-task-paused-label'>
+                  {getLocale(S.CHAT_UI_TASK_STATE_PAUSED_LABEL)}
+                </span>
+              </Label>
+            )}
+
+          {props.isActiveTask
+            && conversationContext.toolUseTaskState
+              === Mojom.TaskState.kStopped && (
+              <Label
+                color='neutral'
+                mode='loud'
+                className={styles.taskStateLabel}
+              >
+                <span data-testid='assistant-task-stopped-label'>
+                  {getLocale(S.CHAT_UI_TASK_STATE_STOPPED_LABEL)}
+                </span>
+              </Label>
+            )}
        </div>
        {taskThumbnail && (
          <div className={styles.taskImage}>
@@ -122,7 +180,7 @@ export default function AssistantTask(props: Props) {
  )
 }

-function Progress(props: Props & { taskData: TaskData }) {
+function Progress(props: Props & TabProps) {
  // The Progress tab should show:
  // - Any last active complete "important" tool (TODO, navigate)
  // - the most recent completion event
@@ -155,6 +213,7 @@ function Progress(props: Props & { taskData: TaskData }) {
          key={index}
          toolUseEvent={event}
          isEntryActive={false}
+          isExecuting={false}
        />
      ))}
      {currentCompletionEvent && (
@@ -173,13 +232,15 @@ function Progress(props: Props & { taskData: TaskData }) {
          key={index}
          toolUseEvent={event.toolUseEvent!}
          isEntryActive={props.isActiveTask}
+          isExecuting={props.isToolExecuting}
        />
      ))}
+      {props.isThinking && <ToolEventThinking />}
    </div>
  )
 }

-function Steps(props: Props & { taskData: TaskData }) {
+function Steps(props: Props & TabProps) {
  // Render every event in the task, split by completion event
  // so that the LLM tells a story of the task by it's own progress
  // description.
@@ -189,7 +250,7 @@ function Steps(props: Props & { taskData: TaskData }) {
      props.isActiveTask && index === props.taskData.taskItems.length - 1

    // Are we generating this task item?
-    const isActive = isRunnable && props.isGenerating && props.isActiveTask
+    const isActive = isRunnable && props.isToolExecuting && props.isActiveTask

    // Are we waiting for this task item to be completed (shouldn't show as complete)
    const isPending =
@@ -199,6 +260,8 @@ function Steps(props: Props & { taskData: TaskData }) {
        (event) => event.toolUseEvent && !event.toolUseEvent.output,
      )

+    const isThinking = isRunnable && props.isThinking
+
    return (
      <div
        key={index}
@@ -222,6 +285,7 @@ function Steps(props: Props & { taskData: TaskData }) {
          allowedLinks={props.taskData.allowedLinks}
          isLeoModel={props.isLeoModel}
        />
+        {isThinking && <ToolEventThinking />}
      </div>
    )
  })
@@ -192,7 +192,6 @@ function ConversationEntries() {
                  <AssistantTask
                    assistantEntries={group}
                    isActiveTask={isLastGroup}
-                    isGenerating={conversationContext.isGenerating}
                    isLeoModel={conversationContext.isLeoModel}
                  />
                )}
@@ -24,6 +24,8 @@ export type ConversationEntriesUIState = Mojom.ConversationEntriesState & {
 export const defaultConversationEntriesUIState: ConversationEntriesUIState = {
  conversationHistory: [],
  isGenerating: false,
+  isToolExecuting: false,
+  toolUseTaskState: Mojom.TaskState.kNone,
  isLeoModel: true,
  allModels: [],
  currentModelKey: '',
@@ -888,4 +888,13 @@
  <message name="IDS_CHAT_UI_STOP_TASK_BUTTON_LABEL" desc="Button label to stop a running task" formatter_data="webui=AiChat">
    Stop
  </message>
+  <message name="IDS_CHAT_UI_TASK_STATE_PAUSED_LABEL" desc="Label for the paused task state" formatter_data="webui=AiChat">
+    Paused
+  </message>
+  <message name="IDS_CHAT_UI_TASK_STATE_STOPPED_LABEL" desc="Label for the stopped task state" formatter_data="webui=AiChat">
+    Stopped
+  </message>
+  <message name="IDS_CHAT_UI_TOOL_LABEL_THINKING" desc="Label for when the assistant is thinking" formatter_data="webui=AiChat">
+    Thinking
+  </message>
 </grit-part>