UI update

azmy60 · azmy60 · commit 468df56a742c · 2025-12-04T20:33:22.000+07:00
diff --git a/src/assets/styles/_base.scss b/src/assets/styles/_base.scss
@@ -1,5 +1,5 @@
 :root {
-  font-family: "Roboto", Helvetica, Arial, sans-serif;
+  font-family: var(--font-family, "Roboto", Helvetica, Arial, sans-serif);
 }
 
 * {
diff --git a/src/components/ChatInterface.vue b/src/components/ChatInterface.vue
@@ -9,10 +9,17 @@
       </div>
       <h1 class="plugin-title">AI Shell</h1>
       <div class="chat-messages">
-        <message v-for="(message, index) in messages" :key="message.id" :message="message"
-          :pending-tool-call-ids="pendingToolCallIds"
-          :status="index === messages.length - 1 ? (status === 'ready' || status === 'error' ? 'ready' : 'processing') : 'ready'"
-          @accept-permission="acceptPermission" @reject-permission="handleRejectPermission" />
+        <template v-for="(message, index) in messages" :key="message.id">
+          <message
+            v-if="
+              !(message.role === 'assistant' && message.parts.find((p) => p.type === 'data-userEditedToolCall'))
+                && !(message.role === 'user' && message.parts.find((p) => p.type === 'data-editedQuery'))
+              "
+            :message="message"
+            :pending-tool-call-ids="pendingToolCallIds"
+            :status="index === messages.length - 1 ? (status === 'ready' || status === 'error' ? 'ready' : 'processing') : 'ready'"
+            @accept-permission="acceptPermission" @reject-permission="handleRejectPermission" />
+        </template>
         <div class="message error" v-if="isUnexpectedError">
           <div class="message-content">
             Something went wrong.
@@ -255,7 +262,7 @@ export default {
 
     handleRejectPermission(options: {
       toolCallId: string;
-      userEditedCode?: string;
+      editedQuery?: string;
     }) {
       this.rejectPermission({
         ...options,
diff --git a/src/components/messages/Message.vue b/src/components/messages/Message.vue
@@ -9,6 +9,7 @@
         </template>
         <tool-message
           v-else-if="isToolUIPart(part)"
+          :message="message"
           :toolCall="part"
           :askingPermission="pendingToolCallIds.includes(part.toolCallId)"
           @accept="$emit('accept-permission', part.toolCallId)"
@@ -37,11 +38,12 @@
 
 <script lang="ts">
 import { PropType } from "vue";
-import { isToolUIPart, UIMessage } from "ai";
+import { isToolUIPart } from "ai";
 import Markdown from "@/components/messages/Markdown.vue";
 import ToolMessage from "@/components/messages/ToolMessage.vue";
 import { clipboard } from "@beekeeperstudio/plugin";
 import { isEmptyUIMessage } from "@/utils";
+import { UIMessage } from "@/types";
 
 export default {
   name: "Message",
diff --git a/src/components/messages/ToolMessage.vue b/src/components/messages/ToolMessage.vue
@@ -1,7 +1,14 @@
 <template>
   <div class="tool" :data-tool-name="name" :data-tool-state="toolCall.state" :data-tool-empty-result="isEmptyResult"
     :data-tool-error="!!error">
-    <div class="tool-name">{{ displayName }}</div>
+    <div class="tool-name">
+      {{ displayName }}
+      <span
+        v-if="message.parts.find((p) => p.type === 'data-toolReplacement')"
+        class="edited-badge"
+        title="You edited the AI's suggestion"
+      >Edited</span>
+    </div>
     <div class="tool-input-container" :style="{ opacity: editing ? 0.5 : 1 }">
       <markdown v-if="name === 'run_query'" :content="'```sql\n' +
         (toolCall.input?.query ||
@@ -21,7 +28,7 @@
         <button class="btn btn-flat" @click="cancelEdit">Cancel</button>
       </div>
     </div>
-    <div v-if="askingPermission && !editing">
+    <div v-if="askingPermission && !editing" class="tool-permission">
       {{
         name === "run_query"
           ? "Do you want to run this query?"
@@ -68,11 +75,16 @@ import { safeJSONStringify } from "@/utils";
 import RunQueryResult from "@/components/messages/tool/RunQueryResult.vue";
 import { isErrorContent, parseErrorContent } from "@/utils";
 import _ from "lodash";
+import { UIMessage } from "@/types";
 
 export default {
   components: { Markdown, RunQueryResult },
   props: {
     askingPermission: Boolean,
+    message: {
+      type: Object as PropType<UIMessage>,
+      required: true,
+    },
     toolCall: {
       type: Object as PropType<ToolUIPart>,
       required: true,
@@ -164,7 +176,7 @@ export default {
     },
     saveEdit() {
       this.editing = false;
-      this.$emit("reject", { userEditedCode: this.queryEditorValue })
+      this.$emit("reject", { editedQuery: this.queryEditorValue })
       this.initialQueryEditorValue = this.queryEditorValue;
     },
     cancelEdit() {
@@ -176,8 +188,8 @@ export default {
 </script>
 
 <style scoped>
-.tool-input-container {
-  margin-bottom: 0.5rem;
+.tool-error {
+  margin-top: 0.5rem;
 }
 
 .tool-input-edit {
@@ -201,6 +213,20 @@ export default {
     gap: 1rem;
     margin-top: 1rem;
   }
+}
+
+.tool-permission {
+  margin-top: 0.5rem;
+}
 
+.edited-badge {
+  border: 1px solid var(--border-color);
+  border-radius: 8px;
+  font-size: 0.75rem;
+  margin-left: 0.25rem;
+  padding-inline: 0.25rem;
+  padding-block: 0.1rem;
+  color: var(--text-light);
+  cursor: default;
 }
 </style>
diff --git a/src/composables/ai.ts b/src/composables/ai.ts
@@ -1,5 +1,5 @@
 import { Chat } from "@ai-sdk/vue";
-import { computed, ComputedRef, watch } from "vue";
+import { computed, ComputedRef, ref, Ref, watch } from "vue";
 import { AvailableProviders, AvailableModels } from "@/config";
 import { tools, userRejectedToolCall } from "@/tools";
 import {
@@ -52,6 +52,8 @@ class AIShellChat {
   readonly pendingToolCallIds: ComputedRef<string[]>;
   readonly askingPermission: ComputedRef<boolean>;
 
+  /** Force the `status` if not `null` */
+  private forceStatus = ref<ChatStatus | null>();
   private chat: Chat<UIMessage>;
   private emitter = mitt<{
     finish: void;
@@ -71,7 +73,7 @@ class AIShellChat {
 
     this.messages = computed(() => this.chat.messages);
     this.error = computed(() => this.chat.error);
-    this.status = computed(() => this.chat.status);
+    this.status = computed(() => this.forceStatus.value ?? this.chat.status);
     this.pendingToolCallIds = computed(() =>
       this.pendingToolCalls.map((t) => t.toolCallId),
     );
@@ -80,11 +82,11 @@ class AIShellChat {
     );
   }
 
-  async send(message: string | UIMessage['parts'], options: SendOptions) {
-    await this.chat.sendMessage(
-      typeof message === "string"
-        ? { text: message }
-        : { parts: message },
+  /** Pass `undefined` to trigger the API without sending the message */
+  async send(message: string | undefined, options: SendOptions) {
+    await this.chat.sendMessage(message
+      ? { text: message }
+      : undefined,
       {
         body: {
           sendOptions: options,
@@ -124,7 +126,7 @@ class AIShellChat {
    **/
   rejectPermission(options?:
     { toolCallId: string; }
-    | { toolCallId: string; userEditedCode: string; sendOptions: SendOptions }
+    | { toolCallId: string; editedQuery: string; sendOptions: SendOptions }
   ) {
     if (!options) {
       this.pendingToolCalls.forEach((t) => {
@@ -138,16 +140,118 @@ class AIShellChat {
       throw new Error(`Tool call with id ${options.toolCallId} not found`);
     }
 
-    if ('userEditedCode' in options) {
+    if ('editedQuery' in options) {
       const sendFollowupMessage = async () => {
         this.emitter.off("finish", sendFollowupMessage);
-        this.send(
-          [{
-            type: "data-userEditedCode",
-            data: { code: options.userEditedCode }
-          }],
-          options.sendOptions
+
+        const assistantMessageId = this.chat.generateId();
+        const replacementToolCallId = this.chat.generateId();
+
+        this.forceStatus.value = "streaming";
+
+        this.chat.messages = [
+          ...this.chat.messages.slice(0, -1),
+          {
+            ...this.chat.lastMessage!,
+            parts: [
+              ...this.chat.lastMessage!.parts,
+              {
+                type: "data-userEditedToolCall",
+                data: { replacementToolCallId },
+              }
+            ],
+          },
+          {
+            id: this.chat.generateId(),
+            role: "user",
+            parts: [{
+              // We use data so it's not shown in the UI
+              type: "data-editedQuery",
+              data: {
+                query: options.editedQuery,
+                targetToolCallId: options.toolCallId,
+              },
+            }],
+          },
+          {
+            id: assistantMessageId,
+            role: "assistant",
+            parts: [
+              { type: "step-start" },
+              {
+                type: "data-toolReplacement",
+                data: { targetToolCallId: options.toolCallId },
+              },
+              {
+                type: "tool-run_query",
+                state: "input-available",
+                toolCallId: replacementToolCallId,
+                input: { query: options.editedQuery },
+              },
+            ],
+          },
+        ];
+
+        const runQueryOutput = await this.createRunQueryToolOutput(
+          replacementToolCallId,
+          options.editedQuery
         );
+
+        if (runQueryOutput.state === "output-error") {
+          this.chat.messages = [
+            ...this.chat.messages.slice(0, -1),
+            {
+              id: assistantMessageId,
+              role: "assistant",
+              parts: [
+                { type: "step-start" },
+                {
+                  type: "data-toolReplacement",
+                  data: { targetToolCallId: options.toolCallId },
+                },
+                {
+                  type: "tool-run_query",
+                  state: "output-error",
+                  toolCallId: replacementToolCallId,
+                  input: { query: options.editedQuery },
+                  errorText: runQueryOutput.errorText,
+                },
+              ],
+            },
+          ];
+        } else {
+          this.chat.messages = [
+            ...this.chat.messages.slice(0, -1),
+            {
+              id: assistantMessageId,
+              role: "assistant",
+              parts: [
+                { type: "step-start" },
+                {
+                  type: "data-toolReplacement",
+                  data: { targetToolCallId: options.toolCallId },
+                },
+                {
+                  type: "tool-run_query",
+                  state: "output-available",
+                  toolCallId: replacementToolCallId,
+                  input: { query: options.editedQuery },
+                  // @ts-expect-error ts doesnt like this because the execute()
+                  // method for run_query (see tools/index.ts) is not defined.
+                  // It can be fixed:
+                  //   1. Upgrading to AI SDK v6
+                  //   2. Use the new API for user permission check
+                  //   3. define execute() method
+                  output: runQueryOutput.output,
+                },
+              ],
+            },
+          ];
+        }
+
+        this.forceStatus.value = null;
+
+        this.send(undefined, options.sendOptions);
       };
       this.emitter.on("finish", sendFollowupMessage);
     }
@@ -264,11 +368,11 @@ class AIShellChat {
       modelId: sendOptions.modelId,
       messages: convertToModelMessages<UIMessage>(m.messages, {
         convertDataPart(part) {
-          if (part.type === "data-userEditedCode") {
+          if (part.type === "data-editedQuery") {
             return {
               type: "text",
               text: "Please run the following code instead:\n```\n"
-                + part.data.code
+                + part.data.query
                 + "\n```",
             };
           }
diff --git a/src/types.ts b/src/types.ts
@@ -2,10 +2,34 @@ import { UIMessage as AIUIMessage, InferUITools } from "ai";
 import { tools } from "./tools";
 
 export type UIDataTypes = {
-  /** The code that the user provided. If defined, it will run the `code` and
-   * give the result to the AI.*/
-  userEditedCode: {
-    code: string;
+  /** Marks the original tool call that was edited by the user.
+   * The UI should show this tool call as replaced/superseded. */
+  userEditedToolCall: {
+    /** The ID of the new tool call that replaces this one. */
+    replacementToolCallId: string;
+  };
+
+  /**
+   * The query or code that the user provided. If defined, AI Shell should run
+   * the `query` and give the result to the AI.
+   *
+   * This data is used by user messages.
+   **/
+  editedQuery: {
+    query: string;
+    /** The tool call that the user rejected. */
+    targetToolCallId: string;
+  };
+
+  /**
+   * If defined, AI Shell should replace the UI of `targetToolCallId` with
+   * this tool or message.
+   *
+   * This data is used by assistant messages.
+   **/
+  toolReplacement: {
+    /** The ID of the original tool call that this replaces. */
+    targetToolCallId: string;
   };
 }
 

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`:root {`
`2`		`- font-family: "Roboto", Helvetica, Arial, sans-serif;`
	`2`	`+ font-family: var(--font-family, "Roboto", Helvetica, Arial, sans-serif);`
`3`	`3`	`}`
`4`	`4`
`5`	`5`	`* {`