diff --git a/cmd/wsh/cmd/wshcmd-ai.go b/cmd/wsh/cmd/wshcmd-ai.go index 15648f297b..b1f4ae3af0 100644 --- a/cmd/wsh/cmd/wshcmd-ai.go +++ b/cmd/wsh/cmd/wshcmd-ai.go @@ -4,47 +4,64 @@ package cmd import ( + "encoding/base64" "fmt" "io" + "net/http" "os" + "path/filepath" "strings" "github.com/spf13/cobra" - "github.com/wavetermdev/waveterm/pkg/waveobj" + "github.com/wavetermdev/waveterm/pkg/util/utilfn" "github.com/wavetermdev/waveterm/pkg/wshrpc" "github.com/wavetermdev/waveterm/pkg/wshrpc/wshclient" "github.com/wavetermdev/waveterm/pkg/wshutil" ) var aiCmd = &cobra.Command{ - Use: "ai [-] [message...]", - Short: "Send a message to an AI block", + Use: "ai [options] [files...]", + Short: "Append content to Wave AI sidebar prompt", + Long: `Append content to Wave AI sidebar prompt (does not auto-submit by default) + +Arguments: + files... Files to attach (use '-' for stdin) + +Examples: + git diff | wsh ai - # Pipe diff to AI, ask question in UI + wsh ai main.go # Attach file, ask question in UI + wsh ai *.go -m "find bugs" # Attach files with message + wsh ai -s - -m "review" < log.txt # Stdin + message, auto-submit + wsh ai -n config.json # New chat with file attached`, RunE: aiRun, PreRunE: preRunSetupRpcClient, DisableFlagsInUseLine: true, } -var aiFileFlags []string +var aiMessageFlag string +var aiSubmitFlag bool var aiNewBlockFlag bool func init() { rootCmd.AddCommand(aiCmd) - aiCmd.Flags().BoolVarP(&aiNewBlockFlag, "new", "n", false, "create a new AI block") - aiCmd.Flags().StringArrayVarP(&aiFileFlags, "file", "f", nil, "attach file content (use '-' for stdin)") + aiCmd.Flags().StringVarP(&aiMessageFlag, "message", "m", "", "optional message/question to append after files") + aiCmd.Flags().BoolVarP(&aiSubmitFlag, "submit", "s", false, "submit the prompt immediately after appending") + aiCmd.Flags().BoolVarP(&aiNewBlockFlag, "new", "n", false, "create a new AI chat instead of using existing") } -func encodeFile(builder *strings.Builder, file io.Reader, fileName string) error { - data, err := io.ReadAll(file) - if err != nil { - return fmt.Errorf("error reading file: %w", err) +func detectMimeType(data []byte) string { + mimeType := http.DetectContentType(data) + return strings.Split(mimeType, ";")[0] +} + +func getMaxFileSize(mimeType string) (int, string) { + if mimeType == "application/pdf" { + return 5 * 1024 * 1024, "5MB" } - // Start delimiter with the file name - builder.WriteString(fmt.Sprintf("\n@@@start file %q\n", fileName)) - // Read the file content and write it to the builder - builder.Write(data) - // End delimiter with the file name - builder.WriteString(fmt.Sprintf("\n@@@end file %q\n\n", fileName)) - return nil + if strings.HasPrefix(mimeType, "image/") { + return 7 * 1024 * 1024, "7MB" + } + return 200 * 1024, "200KB" } func aiRun(cmd *cobra.Command, args []string) (rtnErr error) { @@ -52,118 +69,124 @@ func aiRun(cmd *cobra.Command, args []string) (rtnErr error) { sendActivity("ai", rtnErr == nil) }() - if len(args) == 0 { + if len(args) == 0 && aiMessageFlag == "" { OutputHelpMessage(cmd) - return fmt.Errorf("no message provided") + return fmt.Errorf("no files or message provided") } + const maxFileCount = 15 + const rpcTimeout = 30000 + + var allFiles []wshrpc.AIAttachedFile var stdinUsed bool - var message strings.Builder - // Handle file attachments first - for _, file := range aiFileFlags { - if file == "-" { + if len(args) > maxFileCount { + return fmt.Errorf("too many files (maximum %d files allowed)", maxFileCount) + } + + for _, filePath := range args { + var data []byte + var fileName string + var mimeType string + var err error + + if filePath == "-" { if stdinUsed { return fmt.Errorf("stdin (-) can only be used once") } stdinUsed = true - if err := encodeFile(&message, os.Stdin, ""); err != nil { + + data, err = io.ReadAll(os.Stdin) + if err != nil { return fmt.Errorf("reading from stdin: %w", err) } + fileName = "stdin" + mimeType = "text/plain" } else { - fd, err := os.Open(file) + fileInfo, err := os.Stat(filePath) if err != nil { - return fmt.Errorf("opening file %s: %w", file, err) + return fmt.Errorf("accessing file %s: %w", filePath, err) } - defer fd.Close() - if err := encodeFile(&message, fd, file); err != nil { - return fmt.Errorf("reading file %s: %w", file, err) + if fileInfo.IsDir() { + return fmt.Errorf("%s is a directory, not a file", filePath) } - } - } - // Default to "waveai" block - isDefaultBlock := blockArg == "" - if isDefaultBlock { - blockArg = "view@waveai" - } - var fullORef *waveobj.ORef - var err error - if !aiNewBlockFlag { - fullORef, err = resolveSimpleId(blockArg) - } - if (err != nil && isDefaultBlock) || aiNewBlockFlag { - // Create new AI block if default block doesn't exist - data := &wshrpc.CommandCreateBlockData{ - BlockDef: &waveobj.BlockDef{ - Meta: map[string]interface{}{ - waveobj.MetaKey_View: "waveai", - }, - }, - Focused: true, + data, err = os.ReadFile(filePath) + if err != nil { + return fmt.Errorf("reading file %s: %w", filePath, err) + } + fileName = filepath.Base(filePath) + mimeType = detectMimeType(data) } - newORef, err := wshclient.CreateBlockCommand(RpcClient, *data, &wshrpc.RpcOpts{Timeout: 2000}) - if err != nil { - return fmt.Errorf("creating AI block: %w", err) - } - fullORef = &newORef - // Wait for the block's route to be available - gotRoute, err := wshclient.WaitForRouteCommand(RpcClient, wshrpc.CommandWaitForRouteData{ - RouteId: wshutil.MakeFeBlockRouteId(fullORef.OID), - WaitMs: 4000, - }, &wshrpc.RpcOpts{Timeout: 5000}) - if err != nil { - return fmt.Errorf("waiting for AI block: %w", err) + isPDF := mimeType == "application/pdf" + isImage := strings.HasPrefix(mimeType, "image/") + + if !isPDF && !isImage { + mimeType = "text/plain" + if utilfn.ContainsBinaryData(data) { + return fmt.Errorf("file %s contains binary data and cannot be uploaded as text", fileName) + } } - if !gotRoute { - return fmt.Errorf("AI block route could not be established") + + maxSize, sizeStr := getMaxFileSize(mimeType) + if len(data) > maxSize { + return fmt.Errorf("file %s exceeds maximum size of %s for %s files", fileName, sizeStr, mimeType) } - } else if err != nil { - return fmt.Errorf("resolving block: %w", err) + + allFiles = append(allFiles, wshrpc.AIAttachedFile{ + Name: fileName, + Type: mimeType, + Size: len(data), + Data64: base64.StdEncoding.EncodeToString(data), + }) } - // Create the route for this block - route := wshutil.MakeFeBlockRouteId(fullORef.OID) + tabId := os.Getenv("WAVETERM_TABID") + if tabId == "" { + return fmt.Errorf("WAVETERM_TABID environment variable not set") + } + + route := wshutil.MakeTabRouteId(tabId) - // Then handle main message - if args[0] == "-" { - if stdinUsed { - return fmt.Errorf("stdin (-) can only be used once") + if aiNewBlockFlag { + newChatData := wshrpc.CommandWaveAIAddContextData{ + NewChat: true, } - data, err := io.ReadAll(os.Stdin) + err := wshclient.WaveAIAddContextCommand(RpcClient, newChatData, &wshrpc.RpcOpts{ + Route: route, + Timeout: rpcTimeout, + }) if err != nil { - return fmt.Errorf("reading from stdin: %w", err) - } - message.Write(data) - - // Also include any remaining arguments (excluding the "-" itself) - if len(args) > 1 { - if message.Len() > 0 { - message.WriteString(" ") - } - message.WriteString(strings.Join(args[1:], " ")) + return fmt.Errorf("creating new chat: %w", err) } - } else { - message.WriteString(strings.Join(args, " ")) } - if message.Len() == 0 { - return fmt.Errorf("message is empty") - } - if message.Len() > 50*1024 { - return fmt.Errorf("current max message size is 50k") + for _, file := range allFiles { + contextData := wshrpc.CommandWaveAIAddContextData{ + Files: []wshrpc.AIAttachedFile{file}, + } + err := wshclient.WaveAIAddContextCommand(RpcClient, contextData, &wshrpc.RpcOpts{ + Route: route, + Timeout: rpcTimeout, + }) + if err != nil { + return fmt.Errorf("adding file %s: %w", file.Name, err) + } } - messageData := wshrpc.AiMessageData{ - Message: message.String(), - } - err = wshclient.AiSendMessageCommand(RpcClient, messageData, &wshrpc.RpcOpts{ - Route: route, - Timeout: 2000, - }) - if err != nil { - return fmt.Errorf("sending message: %w", err) + if aiMessageFlag != "" || aiSubmitFlag { + finalContextData := wshrpc.CommandWaveAIAddContextData{ + Text: aiMessageFlag, + Submit: aiSubmitFlag, + } + err := wshclient.WaveAIAddContextCommand(RpcClient, finalContextData, &wshrpc.RpcOpts{ + Route: route, + Timeout: rpcTimeout, + }) + if err != nil { + return fmt.Errorf("adding context: %w", err) + } } return nil diff --git a/docs/docs/wsh-reference.mdx b/docs/docs/wsh-reference.mdx index 0903cde982..d484d9fe3f 100644 --- a/docs/docs/wsh-reference.mdx +++ b/docs/docs/wsh-reference.mdx @@ -110,25 +110,45 @@ wsh getmeta -b [other-tab-id] "bg:*" --clear-prefix | wsh setmeta -b tab --json ## ai -Send messages to new or existing AI blocks directly from the CLI. `-f` passes a file. note that there is a maximum size of 10k for messages and files, so use a tail/grep to cut down file sizes before passing. The `-f` option works great for small files though like shell scripts or `.zshrc` etc. You can use "-" to read input from stdin. +Append content to the Wave AI sidebar. Files are attached as proper file attachments (supporting images, PDFs, and text), not encoded as text. By default, content is added to the sidebar without auto-submitting, allowing you to review and add more context before sending to the AI. -By default the messages get sent to the first AI block (by blocknum). If no AI block exists, then a new one will be created. Use `-n` to force creation of a new AI block. Use `-b` to target a specific AI block. +You can attach multiple files at once (up to 15 files). Use `-m` to add a message along with files, `-s` to auto-submit immediately, and `-n` to start a new chat conversation. Use "-" to read from stdin. ```sh -wsh ai "how do i write an ls command that sorts files in reverse size order" -wsh ai -f <(tail -n 20 "my.log") -- "any idea what these error messages mean" -wsh ai -f README.md "help me update this readme file" +# Pipe command output to AI (ask question in UI) +git diff | wsh ai - +docker logs mycontainer | wsh ai - -# creates a new AI block -wsh ai -n "tell me a story" +# Attach files without auto-submit (review in UI first) +wsh ai main.go utils.go +wsh ai screenshot.png logs.txt -# targets block number 5 -wsh ai -b 5 "tell me more" +# Attach files with message +wsh ai app.py -m "find potential bugs" +wsh ai *.log -m "analyze these error logs" -# read from stdin and also supply a message -tail -n 50 mylog.log | wsh ai - "can you tell me what this error means?" +# Auto-submit immediately +wsh ai config.json -s -m "explain this configuration" +tail -n 50 app.log | wsh ai -s - -m "what's causing these errors?" + +# Start new chat and attach files +wsh ai -n report.pdf data.csv -m "summarize these reports" + +# Attach different file types (images, PDFs, code) +wsh ai architecture.png api-spec.pdf server.go -m "review the system design" ``` +**File Size Limits:** +- Text files: 200KB maximum +- PDF files: 5MB maximum +- Image files: 7MB maximum (accounts for base64 encoding overhead) +- Maximum 15 files per command + +**Flags:** +- `-m, --message ` - Add message text along with files +- `-s, --submit` - Auto-submit immediately (default waits for user) +- `-n, --new` - Clear current chat and start fresh conversation + --- ## editconfig diff --git a/docs/docs/wsh.mdx b/docs/docs/wsh.mdx index ac9b007510..37f5be1fd8 100644 --- a/docs/docs/wsh.mdx +++ b/docs/docs/wsh.mdx @@ -116,17 +116,40 @@ wsh setvar -b tab SHARED_ENV=staging ### AI-Assisted Development +The `wsh ai` command appends content to the Wave AI sidebar. By default, files are attached without auto-submitting, allowing you to review and add more context before sending. + ```bash -# Get AI help with code (uses "-" to read from stdin) -git diff | wsh ai - "review these changes" +# Pipe output to AI sidebar (ask question in UI) +git diff | wsh ai - + +# Attach files with a message +wsh ai main.go utils.go -m "find bugs in these files" + +# Auto-submit with message +wsh ai config.json -s -m "explain this config" -# Get help with a file -wsh ai -f .zshrc "help me add ~/bin to my path" +# Start new chat with attached files +wsh ai -n *.log -m "analyze these logs" -# Debug issues (uses "-" to read from stdin) -dmesg | wsh ai - "help me understand these errors" +# Attach multiple file types (images, PDFs, code) +wsh ai screenshot.png report.pdf app.py -m "review these" + +# Debug with stdin and auto-submit +dmesg | wsh ai -s - -m "help me understand these errors" ``` +**Flags:** +- `-` - Read from stdin instead of a file +- `-m, --message` - Add message text along with files +- `-s, --submit` - Auto-submit immediately (default is to wait for user) +- `-n, --new` - Clear chat and start fresh conversation + +**File Limits:** +- Text files: 200KB max +- PDFs: 5MB max +- Images: 7MB max +- Maximum 15 files per command + ## Tips & Features 1. **Working with Blocks** diff --git a/frontend/app/aipanel/aipanel.tsx b/frontend/app/aipanel/aipanel.tsx index 4231d3add4..aca6f17f93 100644 --- a/frontend/app/aipanel/aipanel.tsx +++ b/frontend/app/aipanel/aipanel.tsx @@ -1,7 +1,6 @@ // Copyright 2025, Command Line Inc. // SPDX-License-Identifier: Apache-2.0 -import { WaveUIMessagePart } from "@/app/aipanel/aitypes"; import { waveAIHasSelection } from "@/app/aipanel/waveai-focus-utils"; import { ErrorBoundary } from "@/app/element/errorboundary"; import { ContextMenuModel } from "@/app/store/contextmenu"; @@ -17,14 +16,14 @@ import { useChat } from "@ai-sdk/react"; import { DefaultChatTransport } from "ai"; import * as jotai from "jotai"; import { memo, useCallback, useEffect, useRef, useState } from "react"; -import { createDataUrl, formatFileSizeError, isAcceptableFile, normalizeMimeType, validateFileSize } from "./ai-utils"; +import { formatFileSizeError, isAcceptableFile, validateFileSize } from "./ai-utils"; import { AIDroppedFiles } from "./aidroppedfiles"; import { AIPanelHeader } from "./aipanelheader"; import { AIPanelInput } from "./aipanelinput"; import { AIPanelMessages } from "./aipanelmessages"; import { AIRateLimitStrip } from "./airatelimitstrip"; import { TelemetryRequiredMessage } from "./telemetryrequired"; -import { WaveAIModel, type DroppedFile } from "./waveai-model"; +import { WaveAIModel } from "./waveai-model"; const AIBlockMask = memo(() => { return ( @@ -195,11 +194,10 @@ interface AIPanelProps { const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { const [isDragOver, setIsDragOver] = useState(false); - const [isLoadingChat, setIsLoadingChat] = useState(true); + const [initialLoadDone, setInitialLoadDone] = useState(false); const model = WaveAIModel.getInstance(); const containerRef = useRef(null); const errorMessage = jotai.useAtomValue(model.errorMessage); - const realMessageRef = useRef(null); const isLayoutMode = jotai.useAtomValue(atoms.controlShiftDelayAtom); const showOverlayBlockNums = jotai.useAtomValue(getSettingsKeyAtom("app:showoverlayblocknums")) ?? true; const focusType = jotai.useAtomValue(focusManager.focusType); @@ -207,12 +205,11 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { const telemetryEnabled = jotai.useAtomValue(getSettingsKeyAtom("telemetry:enabled")) ?? false; const isPanelVisible = jotai.useAtomValue(WorkspaceLayoutModel.getInstance().panelVisibleAtom); - const { messages, sendMessage, status, setMessages, error } = useChat({ + const { messages, sendMessage, status, setMessages, error, stop } = useChat({ transport: new DefaultChatTransport({ api: `${getWebServerEndpoint()}/api/post-chat-message`, prepareSendMessagesRequest: (opts) => { - const msg = realMessageRef.current; - realMessageRef.current = null; + const msg = model.getAndClearMessage(); return { body: { msg, @@ -235,16 +232,17 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { }, }); + model.registerUseChatData(sendMessage, setMessages, status, stop); + // console.log("AICHAT messages", messages); - const clearChat = () => { + const handleClearChat = useCallback(() => { model.clearChat(); - setMessages([]); - }; + }, [model]); const handleKeyDown = (waveEvent: WaveKeyboardEvent): boolean => { if (checkKeyPressed(waveEvent, "Cmd:k")) { - clearChat(); + model.clearChat(); return true; } return false; @@ -259,16 +257,12 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { }, []); useEffect(() => { - const loadMessages = async () => { - const messages = await model.loadChat(); - setMessages(messages as any); - setIsLoadingChat(false); - setTimeout(() => { - model.scrollToBottom(); - }, 100); + const loadChat = async () => { + await model.uiLoadChat(); + setInitialLoadDone(true); }; - loadMessages(); - }, [model, setMessages]); + loadChat(); + }, [model]); useEffect(() => { const updateWidth = () => { @@ -295,69 +289,7 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { const handleSubmit = async (e: React.FormEvent) => { e.preventDefault(); - const input = globalStore.get(model.inputAtom); - if (!input.trim() || status !== "ready" || isLoadingChat) return; - - if (input.trim() === "/clear" || input.trim() === "/new") { - clearChat(); - globalStore.set(model.inputAtom, ""); - return; - } - - model.clearError(); - - const droppedFiles = globalStore.get(model.droppedFiles) as DroppedFile[]; - - // Prepare AI message parts (for backend) - const aiMessageParts: AIMessagePart[] = [{ type: "text", text: input.trim() }]; - - // Prepare UI message parts (for frontend display) - const uiMessageParts: WaveUIMessagePart[] = []; - - if (input.trim()) { - uiMessageParts.push({ type: "text", text: input.trim() }); - } - - // Process files - for (const droppedFile of droppedFiles) { - const normalizedMimeType = normalizeMimeType(droppedFile.file); - const dataUrl = await createDataUrl(droppedFile.file); - - // For AI message (backend) - use data URL - aiMessageParts.push({ - type: "file", - filename: droppedFile.name, - mimetype: normalizedMimeType, - url: dataUrl, - size: droppedFile.file.size, - previewurl: droppedFile.previewUrl, - }); - - uiMessageParts.push({ - type: "data-userfile", - data: { - filename: droppedFile.name, - mimetype: normalizedMimeType, - size: droppedFile.file.size, - previewurl: droppedFile.previewUrl, - }, - }); - } - - // realMessage uses AIMessageParts - const realMessage: AIMessage = { - messageid: crypto.randomUUID(), - parts: aiMessageParts, - }; - realMessageRef.current = realMessage; - - // sendMessage uses UIMessageParts - sendMessage({ parts: uiMessageParts }); - - model.isChatEmpty = false; - globalStore.set(model.inputAtom, ""); - model.clearFiles(); - + await model.handleSubmit(); setTimeout(() => { model.focusInput(); }, 100); @@ -473,7 +405,7 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { menu.push({ label: "New Chat", click: () => { - clearChat(); + model.clearChat(); }, }); @@ -516,7 +448,7 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { > {isDragOver && } {showBlockMask && } - +
@@ -524,7 +456,7 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { ) : ( <> - {messages.length === 0 && !isLoadingChat ? ( + {messages.length === 0 && initialLoadDone ? (
diff --git a/frontend/app/aipanel/aitypes.ts b/frontend/app/aipanel/aitypes.ts index 98bbc0756c..921eb6368f 100644 --- a/frontend/app/aipanel/aitypes.ts +++ b/frontend/app/aipanel/aitypes.ts @@ -1,7 +1,7 @@ // Copyright 2025, Command Line Inc. // SPDX-License-Identifier: Apache-2.0 -import { UIMessage, UIMessagePart } from "ai"; +import { ChatRequestOptions, FileUIPart, UIMessage, UIMessagePart } from "ai"; type WaveUIDataTypes = { userfile: { @@ -23,3 +23,33 @@ type WaveUIDataTypes = { export type WaveUIMessage = UIMessage; export type WaveUIMessagePart = UIMessagePart; + +export type UseChatSetMessagesType = ( + messages: WaveUIMessage[] | ((messages: WaveUIMessage[]) => WaveUIMessage[]) +) => void; + +export type UseChatSendMessageType = ( + message?: + | (Omit & { + id?: string; + role?: "system" | "user" | "assistant"; + } & { + text?: never; + files?: never; + messageId?: string; + }) + | { + text: string; + files?: FileList | FileUIPart[]; + metadata?: unknown; + parts?: never; + messageId?: string; + } + | { + files: FileList | FileUIPart[]; + metadata?: unknown; + parts?: never; + messageId?: string; + }, + options?: ChatRequestOptions +) => Promise; diff --git a/frontend/app/aipanel/waveai-model.tsx b/frontend/app/aipanel/waveai-model.tsx index 8b17c24103..405e0943b1 100644 --- a/frontend/app/aipanel/waveai-model.tsx +++ b/frontend/app/aipanel/waveai-model.tsx @@ -1,15 +1,17 @@ // Copyright 2025, Command Line Inc. // SPDX-License-Identifier: Apache-2.0 +import { UseChatSendMessageType, UseChatSetMessagesType, WaveUIMessagePart } from "@/app/aipanel/aitypes"; import { atoms, getTabMetaKeyAtom } from "@/app/store/global"; import { globalStore } from "@/app/store/jotaiStore"; import * as WOS from "@/app/store/wos"; import { RpcApi } from "@/app/store/wshclientapi"; import { TabRpcClient } from "@/app/store/wshrpcutil"; import { WorkspaceLayoutModel } from "@/app/workspace/workspace-layout-model"; +import { ChatStatus } from "ai"; import * as jotai from "jotai"; import type React from "react"; -import { createImagePreview, resizeImage } from "./ai-utils"; +import { createDataUrl, createImagePreview, normalizeMimeType, resizeImage } from "./ai-utils"; import type { AIPanelInputRef } from "./aipanelinput"; export interface DroppedFile { @@ -25,6 +27,12 @@ export class WaveAIModel { private static instance: WaveAIModel | null = null; private inputRef: React.RefObject | null = null; private scrollToBottomCallback: (() => void) | null = null; + private useChatSendMessage: UseChatSendMessageType | null = null; + private useChatSetMessages: UseChatSetMessagesType | null = null; + private useChatStatus: ChatStatus = "ready"; + private useChatStop: (() => void) | null = null; + // Used for injecting Wave-specific message data into DefaultChatTransport's prepareSendMessagesRequest + realMessage: AIMessage | null = null; widgetAccessAtom!: jotai.Atom; droppedFiles: jotai.PrimitiveAtom = jotai.atom([]); @@ -34,6 +42,7 @@ export class WaveAIModel { containerWidth: jotai.PrimitiveAtom = jotai.atom(0); codeBlockMaxWidth!: jotai.Atom; inputAtom: jotai.PrimitiveAtom = jotai.atom(""); + isLoadingChatAtom: jotai.PrimitiveAtom = jotai.atom(false); isChatEmpty: boolean = true; private constructor() { @@ -127,6 +136,7 @@ export class WaveAIModel { } clearChat() { + this.useChatStop?.(); this.clearFiles(); this.isChatEmpty = true; const newChatId = crypto.randomUUID(); @@ -137,6 +147,8 @@ export class WaveAIModel { oref: WOS.makeORef("tab", tabId), meta: { "waveai:chatid": newChatId }, }); + + this.useChatSetMessages?.([]); } setError(message: string) { @@ -155,6 +167,18 @@ export class WaveAIModel { this.scrollToBottomCallback = callback; } + registerUseChatData( + sendMessage: UseChatSendMessageType, + setMessages: UseChatSetMessagesType, + status: ChatStatus, + stop: () => void + ) { + this.useChatSendMessage = sendMessage; + this.useChatSetMessages = setMessages; + this.useChatStatus = status; + this.useChatStop = stop; + } + scrollToBottom() { this.scrollToBottomCallback?.(); } @@ -168,11 +192,29 @@ export class WaveAIModel { } } + getAndClearMessage(): AIMessage | null { + const msg = this.realMessage; + this.realMessage = null; + return msg; + } + hasNonEmptyInput(): boolean { const input = globalStore.get(this.inputAtom); return input != null && input.trim().length > 0; } + appendText(text: string) { + const currentInput = globalStore.get(this.inputAtom); + let newInput = currentInput; + + if (newInput.length > 0 && !newInput.endsWith(" ") && !newInput.endsWith("\n")) { + newInput += " "; + } + + newInput += text; + globalStore.set(this.inputAtom, newInput); + } + setModel(model: string) { const tabId = globalStore.get(atoms.staticTabId); RpcApi.SetMetaCommand(TabRpcClient, { @@ -214,6 +256,83 @@ export class WaveAIModel { } } + async handleSubmit() { + const input = globalStore.get(this.inputAtom); + const droppedFiles = globalStore.get(this.droppedFiles); + + if (input.trim() === "/clear" || input.trim() === "/new") { + this.clearChat(); + globalStore.set(this.inputAtom, ""); + return; + } + + if ( + (!input.trim() && droppedFiles.length === 0) || + (this.useChatStatus !== "ready" && this.useChatStatus !== "error") || + globalStore.get(this.isLoadingChatAtom) + ) { + return; + } + + this.clearError(); + + const aiMessageParts: AIMessagePart[] = []; + const uiMessageParts: WaveUIMessagePart[] = []; + + if (input.trim()) { + aiMessageParts.push({ type: "text", text: input.trim() }); + uiMessageParts.push({ type: "text", text: input.trim() }); + } + + for (const droppedFile of droppedFiles) { + const normalizedMimeType = normalizeMimeType(droppedFile.file); + const dataUrl = await createDataUrl(droppedFile.file); + + aiMessageParts.push({ + type: "file", + filename: droppedFile.name, + mimetype: normalizedMimeType, + url: dataUrl, + size: droppedFile.file.size, + previewurl: droppedFile.previewUrl, + }); + + uiMessageParts.push({ + type: "data-userfile", + data: { + filename: droppedFile.name, + mimetype: normalizedMimeType, + size: droppedFile.file.size, + previewurl: droppedFile.previewUrl, + }, + }); + } + + const realMessage: AIMessage = { + messageid: crypto.randomUUID(), + parts: aiMessageParts, + }; + this.realMessage = realMessage; + + // console.log("SUBMIT MESSAGE", realMessage); + + this.useChatSendMessage?.({ parts: uiMessageParts }); + + this.isChatEmpty = false; + globalStore.set(this.inputAtom, ""); + this.clearFiles(); + } + + async uiLoadChat() { + globalStore.set(this.isLoadingChatAtom, true); + const messages = await this.loadChat(); + this.useChatSetMessages?.(messages as any); + globalStore.set(this.isLoadingChatAtom, false); + setTimeout(() => { + this.scrollToBottom(); + }, 100); + } + async ensureRateLimitSet() { const currentInfo = globalStore.get(atoms.waveAIRateLimitInfoAtom); if (currentInfo != null) { diff --git a/frontend/app/store/tabrpcclient.ts b/frontend/app/store/tabrpcclient.ts index b74d7e8120..f0c6da55ce 100644 --- a/frontend/app/store/tabrpcclient.ts +++ b/frontend/app/store/tabrpcclient.ts @@ -1,8 +1,11 @@ // Copyright 2025, Command Line Inc. // SPDX-License-Identifier: Apache-2.0 +import { WaveAIModel } from "@/app/aipanel/waveai-model"; import { getApi } from "@/app/store/global"; +import { WorkspaceLayoutModel } from "@/app/workspace/workspace-layout-model"; import { getLayoutModelForStaticTab } from "@/layout/index"; +import { base64ToArray } from "@/util/util"; import { RpcResponseHelper, WshClient } from "./wshclient"; export class TabClient extends WshClient { @@ -56,4 +59,34 @@ export class TabClient extends WshClient { return await getApi().captureScreenshot(electronRect); } + + async handle_waveaiaddcontext(rh: RpcResponseHelper, data: CommandWaveAIAddContextData): Promise { + const workspaceLayoutModel = WorkspaceLayoutModel.getInstance(); + if (!workspaceLayoutModel.getAIPanelVisible()) { + workspaceLayoutModel.setAIPanelVisible(true, { nofocus: true }); + } + + const model = WaveAIModel.getInstance(); + + if (data.newchat) { + model.clearChat(); + } + + if (data.files && data.files.length > 0) { + for (const fileData of data.files) { + const decodedData = base64ToArray(fileData.data64); + const blob = new Blob([decodedData], { type: fileData.type }); + const file = new File([blob], fileData.name, { type: fileData.type }); + await model.addFile(file); + } + } + + if (data.text) { + model.appendText(data.text); + } + + if (data.submit) { + await model.handleSubmit(); + } + } } diff --git a/frontend/app/store/wshclientapi.ts b/frontend/app/store/wshclientapi.ts index dc81b159ac..e2220ea827 100644 --- a/frontend/app/store/wshclientapi.ts +++ b/frontend/app/store/wshclientapi.ts @@ -487,6 +487,11 @@ class RpcApiType { return client.wshRpcCall("waitforroute", data, opts); } + // command "waveaiaddcontext" [call] + WaveAIAddContextCommand(client: WshClient, data: CommandWaveAIAddContextData, opts?: RpcOpts): Promise { + return client.wshRpcCall("waveaiaddcontext", data, opts); + } + // command "waveaienabletelemetry" [call] WaveAIEnableTelemetryCommand(client: WshClient, opts?: RpcOpts): Promise { return client.wshRpcCall("waveaienabletelemetry", null, opts); diff --git a/frontend/app/workspace/workspace-layout-model.ts b/frontend/app/workspace/workspace-layout-model.ts index da1ffa4a1c..9b5437316f 100644 --- a/frontend/app/workspace/workspace-layout-model.ts +++ b/frontend/app/workspace/workspace-layout-model.ts @@ -218,7 +218,7 @@ class WorkspaceLayoutModel { return this.aiPanelVisible; } - setAIPanelVisible(visible: boolean): void { + setAIPanelVisible(visible: boolean, opts?: { nofocus?: boolean }): void { if (this.focusTimeoutRef != null) { clearTimeout(this.focusTimeoutRef); this.focusTimeoutRef = null; @@ -238,10 +238,12 @@ class WorkspaceLayoutModel { this.syncAIPanelRef(); if (visible) { - this.focusTimeoutRef = setTimeout(() => { - WaveAIModel.getInstance().focusInput(); - this.focusTimeoutRef = null; - }, 350); + if (!opts?.nofocus) { + this.focusTimeoutRef = setTimeout(() => { + WaveAIModel.getInstance().focusInput(); + this.focusTimeoutRef = null; + }, 350); + } } else { const layoutModel = getLayoutModelForStaticTab(); const focusedNode = globalStore.get(layoutModel.focusedNode); diff --git a/frontend/types/gotypes.d.ts b/frontend/types/gotypes.d.ts index fd0cfd0a4e..3b88370c4a 100644 --- a/frontend/types/gotypes.d.ts +++ b/frontend/types/gotypes.d.ts @@ -5,6 +5,14 @@ declare global { + // wshrpc.AIAttachedFile + type AIAttachedFile = { + name: string; + type: string; + size: number; + data64: string; + }; + // wshrpc.ActivityDisplayType type ActivityDisplayType = { width: number; @@ -320,6 +328,14 @@ declare global { waitms: number; }; + // wshrpc.CommandWaveAIAddContextData + type CommandWaveAIAddContextData = { + files?: AIAttachedFile[]; + text?: string; + submit?: boolean; + newchat?: boolean; + }; + // wshrpc.CommandWaveAIToolApproveData type CommandWaveAIToolApproveData = { toolcallid: string; diff --git a/frontend/util/util.ts b/frontend/util/util.ts index 4a081ad64d..69e17fceb2 100644 --- a/frontend/util/util.ts +++ b/frontend/util/util.ts @@ -4,8 +4,8 @@ import base64 from "base64-js"; import clsx, { type ClassValue } from "clsx"; import { Atom, atom, Getter, SetStateAction, Setter, useAtomValue } from "jotai"; -import { debounce, throttle } from "throttle-debounce"; import { twMerge } from "tailwind-merge"; +import { debounce, throttle } from "throttle-debounce"; const prevValueCache = new WeakMap(); // stores a previous value for a deep equal comparison (used with the deepCompareReturnPrev function) function isBlank(str: string): boolean { @@ -28,7 +28,7 @@ function stringToBase64(input: string): string { return base64.fromByteArray(stringBytes); } -function base64ToArray(b64: string): Uint8Array { +function base64ToArray(b64: string): Uint8Array { const rawStr = atob(b64); const rtnArr = new Uint8Array(new ArrayBuffer(rawStr.length)); for (let i = 0; i < rawStr.length; i++) { @@ -379,17 +379,22 @@ function mergeMeta(meta: MetaType, metaUpdate: MetaType, prefix?: string): MetaT } function escapeBytes(str: string): string { - return str.replace(/[\s\S]/g, ch => { + return str.replace(/[\s\S]/g, (ch) => { const code = ch.charCodeAt(0); switch (ch) { - case "\n": return "\\n"; - case "\r": return "\\r"; - case "\t": return "\\t"; - case "\b": return "\\b"; - case "\f": return "\\f"; + case "\n": + return "\\n"; + case "\r": + return "\\r"; + case "\t": + return "\\t"; + case "\b": + return "\\b"; + case "\f": + return "\\f"; } if (code === 0x1b) return "\\x1b"; // escape - if (code < 0x20 || code === 0x7f) return `\\x${code.toString(16).padStart(2,"0")}`; + if (code < 0x20 || code === 0x7f) return `\\x${code.toString(16).padStart(2, "0")}`; return ch; }); } diff --git a/pkg/aiusechat/openai/openai-convertmessage.go b/pkg/aiusechat/openai/openai-convertmessage.go index 4fccea00a0..8c9559d699 100644 --- a/pkg/aiusechat/openai/openai-convertmessage.go +++ b/pkg/aiusechat/openai/openai-convertmessage.go @@ -12,10 +12,12 @@ import ( "fmt" "log" "net/http" + "strconv" "strings" "github.com/google/uuid" "github.com/wavetermdev/waveterm/pkg/aiusechat/uctypes" + "github.com/wavetermdev/waveterm/pkg/util/utilfn" ) const ( @@ -23,6 +25,38 @@ const ( OpenAIDefaultMaxTokens = 4096 ) +// extractXmlAttribute extracts an attribute value from an XML-like tag. +// Expects double-quoted strings where internal quotes are encoded as ". +// Returns the unquoted value and true if found, or empty string and false if not found or invalid. +func extractXmlAttribute(tag, attrName string) (string, bool) { + attrStart := strings.Index(tag, attrName+"=") + if attrStart == -1 { + return "", false + } + + pos := attrStart + len(attrName+"=") + start := strings.Index(tag[pos:], `"`) + if start == -1 { + return "", false + } + start += pos + + end := strings.Index(tag[start+1:], `"`) + if end == -1 { + return "", false + } + end += start + 1 + + quotedValue := tag[start : end+1] + value, err := strconv.Unquote(quotedValue) + if err != nil { + return "", false + } + + value = strings.ReplaceAll(value, """, `"`) + return value, true +} + // ---------- OpenAI Request Types ---------- type StreamOptionsType struct { @@ -292,24 +326,34 @@ func convertFileAIMessagePart(part uctypes.AIMessagePart) (*OpenAIMessageContent }, nil case part.MimeType == "text/plain": - // Handle text/plain files as input_text with special formatting var textContent string if len(part.Data) > 0 { textContent = string(part.Data) } else if part.URL != "" { - return nil, fmt.Errorf("dropping text/plain file with URL (must be fetched and converted to data)") + if strings.HasPrefix(part.URL, "data:") { + _, decodedData, err := utilfn.DecodeDataURL(part.URL) + if err != nil { + return nil, fmt.Errorf("failed to decode data URL for text/plain file: %w", err) + } + textContent = string(decodedData) + } else { + return nil, fmt.Errorf("dropping text/plain file with URL (must be fetched and converted to data)") + } } else { return nil, fmt.Errorf("text/plain file part missing data") } - // Format as: file "filename" (mimetype)\n\nfile-content fileName := part.FileName if fileName == "" { fileName = "untitled.txt" } - formattedText := fmt.Sprintf("file %q (%s)\n\n%s", fileName, part.MimeType, textContent) + encodedFileName := strings.ReplaceAll(fileName, `"`, """) + quotedFileName := strconv.Quote(encodedFileName) + + randomSuffix := uuid.New().String()[0:8] + formattedText := fmt.Sprintf("\n%s\n", randomSuffix, quotedFileName, textContent, randomSuffix) return &OpenAIMessageContent{ Type: "input_text", @@ -435,11 +479,31 @@ func (m *OpenAIChatMessage) ConvertToUIMessage() *uctypes.UIMessage { for _, block := range m.Message.Content { switch block.Type { case "input_text", "output_text": - // Convert text blocks to UIMessagePart - parts = append(parts, uctypes.UIMessagePart{ - Type: "text", - Text: block.Text, - }) + if strings.HasPrefix(block.Text, "\ncontent\n.`, + `If multiple attached files exist, treat each as a separate source file with its own file_name.`, + `When the user refers to these files, use their inline content directly; do NOT call any read_text_file or file-access tools to re-read them unless asked.`, + // Output & formatting `When presenting commands or any runnable multi-line code, always use fenced Markdown code blocks.`, `Use an appropriate language hint after the opening fence (e.g., "bash" for shell commands, "go" for Go, "json" for JSON).`, diff --git a/pkg/util/utilfn/marshal.go b/pkg/util/utilfn/marshal.go index 63c4617895..c284260ae9 100644 --- a/pkg/util/utilfn/marshal.go +++ b/pkg/util/utilfn/marshal.go @@ -5,8 +5,10 @@ package utilfn import ( "bytes" + "encoding/base64" "encoding/json" "fmt" + "net/url" "reflect" "strings" @@ -162,3 +164,54 @@ func setValue(field reflect.Value, value any) error { return fmt.Errorf("cannot set value of type %v to field of type %v", valueRef.Type(), field.Type()) } + +// DecodeDataURL decodes a data URL and returns the mimetype and raw data bytes +func DecodeDataURL(dataURL string) (mimeType string, data []byte, err error) { + if !strings.HasPrefix(dataURL, "data:") { + return "", nil, fmt.Errorf("invalid data URL: must start with 'data:'") + } + + parts := strings.SplitN(dataURL, ",", 2) + if len(parts) != 2 { + return "", nil, fmt.Errorf("invalid data URL format: missing comma separator") + } + + header := parts[0] + dataStr := parts[1] + + // Parse mimetype from header: "data:text/plain;base64" -> "text/plain" + headerWithoutPrefix := strings.TrimPrefix(header, "data:") + mimeType = strings.Split(headerWithoutPrefix, ";")[0] + if mimeType == "" { + mimeType = "text/plain" // default mimetype + } + + if strings.Contains(header, ";base64") { + decoded, decodeErr := base64.StdEncoding.DecodeString(dataStr) + if decodeErr != nil { + return "", nil, fmt.Errorf("failed to decode base64 data: %w", decodeErr) + } + return mimeType, decoded, nil + } + + // Non-base64 data URLs are percent-encoded + decoded, decodeErr := url.QueryUnescape(dataStr) + if decodeErr != nil { + return "", nil, fmt.Errorf("failed to decode percent-encoded data: %w", decodeErr) + } + return mimeType, []byte(decoded), nil +} + + +// ContainsBinaryData checks if the provided data contains binary (non-text) content +func ContainsBinaryData(data []byte) bool { + for _, b := range data { + if b == 0 { + return true + } + if b < 32 && b != 9 && b != 10 && b != 13 { + return true + } + } + return false +} diff --git a/pkg/web/ws.go b/pkg/web/ws.go index f33eee9f84..fcf110556d 100644 --- a/pkg/web/ws.go +++ b/pkg/web/ws.go @@ -27,7 +27,7 @@ const wsReadWaitTimeout = 15 * time.Second const wsWriteWaitTimeout = 10 * time.Second const wsPingPeriodTickTime = 10 * time.Second const wsInitialPingTime = 1 * time.Second -const wsMaxMessageSize = 8 * 1024 * 1024 +const wsMaxMessageSize = 10 * 1024 * 1024 const DefaultCommandTimeout = 2 * time.Second diff --git a/pkg/wshrpc/wshclient/wshclient.go b/pkg/wshrpc/wshclient/wshclient.go index 02d0c7242b..47af277974 100644 --- a/pkg/wshrpc/wshclient/wshclient.go +++ b/pkg/wshrpc/wshclient/wshclient.go @@ -581,6 +581,12 @@ func WaitForRouteCommand(w *wshutil.WshRpc, data wshrpc.CommandWaitForRouteData, return resp, err } +// command "waveaiaddcontext", wshserver.WaveAIAddContextCommand +func WaveAIAddContextCommand(w *wshutil.WshRpc, data wshrpc.CommandWaveAIAddContextData, opts *wshrpc.RpcOpts) error { + _, err := sendRpcRequestCallHelper[any](w, "waveaiaddcontext", data, opts) + return err +} + // command "waveaienabletelemetry", wshserver.WaveAIEnableTelemetryCommand func WaveAIEnableTelemetryCommand(w *wshutil.WshRpc, opts *wshrpc.RpcOpts) error { _, err := sendRpcRequestCallHelper[any](w, "waveaienabletelemetry", nil, opts) diff --git a/pkg/wshrpc/wshrpctypes.go b/pkg/wshrpc/wshrpctypes.go index 754b0fa4cf..f1abf21e04 100644 --- a/pkg/wshrpc/wshrpctypes.go +++ b/pkg/wshrpc/wshrpctypes.go @@ -143,6 +143,7 @@ const ( Command_GetWaveAIChat = "getwaveaichat" Command_GetWaveAIRateLimit = "getwaveairatelimit" Command_WaveAIToolApprove = "waveaitoolapprove" + Command_WaveAIAddContext = "waveaiaddcontext" Command_CaptureBlockScreenshot = "captureblockscreenshot" @@ -273,6 +274,7 @@ type WshRpcInterface interface { GetWaveAIChatCommand(ctx context.Context, data CommandGetWaveAIChatData) (*uctypes.UIChat, error) GetWaveAIRateLimitCommand(ctx context.Context) (*uctypes.RateLimitInfo, error) WaveAIToolApproveCommand(ctx context.Context, data CommandWaveAIToolApproveData) error + WaveAIAddContextCommand(ctx context.Context, data CommandWaveAIAddContextData) error // screenshot CaptureBlockScreenshotCommand(ctx context.Context, data CommandCaptureBlockScreenshotData) (string, error) @@ -734,6 +736,20 @@ type CommandWaveAIToolApproveData struct { Approval string `json:"approval,omitempty"` } +type AIAttachedFile struct { + Name string `json:"name"` + Type string `json:"type"` + Size int `json:"size"` + Data64 string `json:"data64"` +} + +type CommandWaveAIAddContextData struct { + Files []AIAttachedFile `json:"files,omitempty"` + Text string `json:"text,omitempty"` + Submit bool `json:"submit,omitempty"` + NewChat bool `json:"newchat,omitempty"` +} + type CommandCaptureBlockScreenshotData struct { BlockId string `json:"blockid" wshcontext:"BlockId"` }