feat: enhance LM Studio stability and session management v2.2.27

src/agent.ts

@@ -1,13 +1,12 @@
 import * as vscode from 'vscode';
 import * as path from 'path';
 import * as fs from 'fs';
-import axios from 'axios';
+// axios removed
 import { 
     getConfig, 
     _getBrainDir, 
+    findBrainFiles,
     EXCLUDED_DIRS, 
-    MAX_CONTEXT_SIZE, 
-    MAX_AUTO_AGENT_STEPS,
     SYSTEM_PROMPT,
     shouldAutoPushBrain,
     getSecondBrainRepo
@@ -20,13 +19,37 @@ export interface ChatMessage {
 }
 
 export class AgentExecutor {
+    private chatHistory: ChatMessage[] = [];
+    private abortController: AbortController | null = null;
+    private webview: vscode.Webview | undefined;
+
     constructor(
-        private context: vscode.ExtensionContext,
-        private webview: vscode.Webview | undefined,
-        private chatHistory: ChatMessage[],
-        private abortController: AbortController | null
+        private context: vscode.ExtensionContext
     ) {}
 
+    public setWebview(webview: vscode.Webview) {
+        this.webview = webview;
+    }
+
+    public getHistory() {
+        return this.chatHistory;
+    }
+
+    public setHistory(history: ChatMessage[]) {
+        this.chatHistory = history;
+    }
+
+    public clearHistory() {
+        this.chatHistory = [];
+    }
+
+    public stop() {
+        if (this.abortController) {
+            this.abortController.abort();
+            this.abortController = null;
+        }
+    }
+
     public async handlePrompt(
         prompt: string | null,
         modelName: string,
@@ -56,11 +79,12 @@ export class AgentExecutor {
             const rootPath = workspaceFolders ? workspaceFolders[0].uri.fsPath : '';
             
             let contextBlock = '';
+            const config = getConfig();
             const editor = vscode.window.activeTextEditor;
             if (editor && editor.document.uri.scheme === 'file') {
                 const text = editor.document.getText();
                 const name = path.basename(editor.document.fileName);
-                if (text.trim().length > 0 && text.length < MAX_CONTEXT_SIZE) {
+                if (text.trim().length > 0 && text.length < config.maxContextSize) {
                     contextBlock = `\n\n[Currently open file: ${name}]\n\`\`\`\n${text}\n\`\`\``;
                 }
             }
@@ -96,15 +120,21 @@ export class AgentExecutor {
                     if (typeof content === 'string') {
                         reqMessages[firstUserIdx].content = `${fullSystemPrompt}\n\n[USER QUERY]\n${content}`;
                         if (loopDepth > 0) {
-                            reqMessages[firstUserIdx].content = `[Autonomous Step ${loopDepth}/${MAX_AUTO_AGENT_STEPS}]\n${reqMessages[firstUserIdx].content}`;
+                            reqMessages[firstUserIdx].content = `[Autonomous Step ${loopDepth}/${config.maxAutoSteps}]\n${reqMessages[firstUserIdx].content}`;
                         }
                     }
                 }
             }
 
             // 4. Call AI Engine
-            const isLMStudio = ollamaUrl.includes('1234') || ollamaUrl.includes('v1');
-            const apiUrl = isLMStudio ? `${ollamaUrl}/v1/chat/completions` : `${ollamaUrl}/api/chat`;
+            const isLMStudio = ollamaUrl.includes('1234') || ollamaUrl.includes('v1') || ollamaUrl.includes('localhost'); 
+            // Note: Many users use LM Studio on localhost, we'll try to be smart or fallback to Ollama format if it fails.
+            
+            const apiUrl = isLMStudio ? 
+                (ollamaUrl.endsWith('/v1') ? `${ollamaUrl}/chat/completions` : `${ollamaUrl}/v1/chat/completions`) : 
+                `${ollamaUrl}/api/chat`;
+
+            this.abortController = new AbortController();
 
             const streamBody = {
                 model: modelName || defaultModel,
@@ -112,41 +142,79 @@ export class AgentExecutor {
                 stream: true,
                 ...(isLMStudio 
                     ? { max_tokens: 4096, temperature } 
-                    : { options: { num_ctx: 16384, num_predict: 4096, temperature } }),
+                    : { options: { num_ctx: 32768, num_predict: 4096, temperature } }),
             };
 
+            const response = await fetch(apiUrl, {
+                method: 'POST',
+                headers: { 
+                    'Content-Type': 'application/json',
+                    'Accept': 'text/event-stream',
+                    'Cache-Control': 'no-cache',
+                    'Connection': 'keep-alive'
+                },
+                body: JSON.stringify(streamBody),
+                signal: this.abortController.signal,
+                keepalive: true
+            });
+
+            if (!response.ok) {
+                const errText = await response.text();
+                throw new Error(`AI Engine error: ${response.status} - ${errText}`);
+            }
+
+            let aiResponseText = '';
+            const reader = response.body?.getReader();
+            if (!reader) throw new Error("Response body is not readable.");
+
             if (loopDepth === 0) this.webview.postMessage({ type: 'streamStart' });
 
-            const response = await axios.post(apiUrl, streamBody, {
-                timeout,
-                responseType: 'stream',
-                signal: this.abortController?.signal
-            });
+            let buffer = '';
+            const decoder = new TextDecoder();
+            try {
+                while (true) {
+                    const { done, value } = await reader.read();
+                    if (done) break;
 
-            let aiResponseText = '';
-            await new Promise<void>((resolve, reject) => {
-                const stream = response.data;
-                let buffer = '';
-                stream.on('data', (chunk: Buffer) => {
-                    buffer += chunk.toString();
+                    buffer += decoder.decode(value, { stream: true });
                     const lines = buffer.split('\n');
                     buffer = lines.pop() || '';
                     for (const line of lines) {
-                        if (!line.trim() || line.trim() === 'data: [DONE]') continue;
+                        const trimmed = line.trim();
+                        if (!trimmed || trimmed === 'data: [DONE]') continue;
                         try {
-                            const raw = line.startsWith('data: ') ? line.slice(6) : line;
+                            const raw = trimmed.startsWith('data: ') ? trimmed.slice(6) : trimmed;
                             const json = JSON.parse(raw);
-                            const token = isLMStudio ? json.choices?.[0]?.delta?.content || '' : json.message?.content || '';
+                            const token = isLMStudio ? json.choices?.[0]?.delta?.content || '' : json.message?.content || json.response || '';
                             if (token) {
                                 aiResponseText += token;
                                 this.webview?.postMessage({ type: 'streamChunk', value: token });
                             }
-                        } catch {}
+                        } catch (e) {}
                     }
-                });
-                stream.on('end', () => resolve());
-                stream.on('error', (err: any) => reject(err));
-            });
+                }
+            } catch (err: any) {
+                if (err.name === 'AbortError') {
+                    console.log('[Agent] Generation aborted by user.');
+                } else {
+                    console.error('[Agent] Stream reading error:', err);
+                    this.webview?.postMessage({ type: 'error', value: `Connection lost: ${err.message}` });
+                }
+            }
+
+            // Final buffer processing
+            if (buffer.trim() && buffer.trim() !== 'data: [DONE]') {
+                try {
+                    const trimmed = buffer.trim();
+                    const raw = trimmed.startsWith('data: ') ? trimmed.slice(6) : trimmed;
+                    const json = JSON.parse(raw);
+                    const token = isLMStudio ? json.choices?.[0]?.delta?.content || '' : json.message?.content || json.response || '';
+                    if (token) {
+                        aiResponseText += token;
+                        this.webview?.postMessage({ type: 'streamChunk', value: token });
+                    }
+                } catch (e) {}
+            }
 
             if (loopDepth === 0) this.webview.postMessage({ type: 'streamEnd' });
             this.chatHistory.push({ role: 'assistant', content: aiResponseText });
@@ -155,24 +223,28 @@ export class AgentExecutor {
             const report = await this.executeActions(aiResponseText, rootPath);
 
             if (report.length > 0) {
-                const reportMsg = `\n\n---\n**[Agent Action Report] (${loopDepth + 1}/${MAX_AUTO_AGENT_STEPS})**\n${report.join("\n")}`;
+                const reportMsg = `\n\n> ⚙️ **System Action Report** (${loopDepth + 1}/${config.maxAutoSteps})\n> ${report.join("\n> ")}\n\n`;
                 this.webview.postMessage({ type: 'streamChunk', value: reportMsg });
                 
                 // Continue loop if needed
-                if (loopDepth < MAX_AUTO_AGENT_STEPS) {
+                if (loopDepth < config.maxAutoSteps) {
                     const currentActionStr = report.join('|');
                     const lastActionStr = this.context.workspaceState.get<string>('lastActionStr');
                     
                     if (currentActionStr === lastActionStr) {
-                        this.webview.postMessage({ type: "streamChunk", value: "\n\n**[Loop Detected]** AI is repeating actions. Diverging..." });
-                        this.chatHistory.push({ role: "user", content: "[System] Action repeated. Try a different strategy." });
+                        this.webview.postMessage({ type: 'streamChunk', value: "\n⚠️ *Stopping to prevent infinite loop.*" });
+                        if (loopDepth === 0) this.webview.postMessage({ type: 'streamEnd' });
+                        return;
                     }
                     
                     await this.context.workspaceState.update('lastActionStr', currentActionStr);
-                    this.webview.postMessage({ type: 'autoContinue', value: `Thinking... (${loopDepth + 1}/${MAX_AUTO_AGENT_STEPS})` });
                     
+                    // Explicitly tell the AI to look at the results and continue
+                    const continuationPrompt = "I have executed your actions. Above is the result. Please analyze it and provide the next step or the final answer.";
+                    
+                    this.webview.postMessage({ type: 'autoContinue', value: `Thinking... (${loopDepth + 1}/${config.maxAutoSteps})` });
                     await new Promise(r => setTimeout(r, 800));
-                    await this.handlePrompt(null, modelName, { ...options, loopDepth: loopDepth + 1 });
+                    await this.handlePrompt(continuationPrompt, modelName, { ...options, loopDepth: loopDepth + 1 });
                 }
             }
 
@@ -233,7 +305,22 @@ export class AgentExecutor {
             } catch (err: any) { report.push(`❌ Error Editing ${relPath}: ${err.message}`); }
         }
 
-        // Action 3: Read File
+        // Action 3: Delete File
+        const deleteRegex = /<delete_file\s+path=['"]?([^'"]+)['"]?\s*\/?>(?:<\/delete_file>)?/gi;
+        while ((match = deleteRegex.exec(aiMessage)) !== null) {
+            const relPath = match[1].trim();
+            try {
+                const absPath = validatePath(rootPath, relPath);
+                if (fs.existsSync(absPath)) {
+                    fs.unlinkSync(absPath);
+                    report.push(`🗑 Deleted: ${relPath}`);
+                } else {
+                    report.push(`⚠️ Delete failed: ${relPath} not found`);
+                }
+            } catch (err: any) { report.push(`❌ Error Deleting ${relPath}: ${err.message}`); }
+        }
+
+        // Action 4: Read File
         const readRegex = /<read_file\s+path=['"]?([^'"]+)['"]?\s*\/?>(?:<\/read_file>)?/gi;
         while ((match = readRegex.exec(aiMessage)) !== null) {
             const relPath = match[1].trim();
@@ -250,7 +337,7 @@ export class AgentExecutor {
             } catch (err: any) { report.push(`❌ Error Reading ${relPath}: ${err.message}`); }
         }
 
-        // Action 4: Run Command
+        // Action 5: Run Command
         const cmdRegex = /<run_command>([\s\S]*?)<\/run_command>/gi;
         while ((match = cmdRegex.exec(aiMessage)) !== null) {
             const cmd = match[1].trim();
@@ -263,7 +350,7 @@ export class AgentExecutor {
             } catch (err: any) { report.push(`❌ Blocked: ${err.message}`); }
         }
 
-        // Action 5: List Files
+        // Action 6: List Files
         const listRegex = /<list_files\s+path=['"]?([^'"]+)['"]?\s*\/?>(?:<\/list_files>)?/gi;
         while ((match = listRegex.exec(aiMessage)) !== null) {
             const relPath = match[1].trim() || '.';
@@ -271,16 +358,81 @@ export class AgentExecutor {
                 const absPath = validatePath(rootPath, relPath);
                 if (fs.existsSync(absPath) && fs.statSync(absPath).isDirectory()) {
                     const entries = fs.readdirSync(absPath, { withFileTypes: true });
-                    const listing = entries
+                    let listing = entries
                         .filter(e => !e.name.startsWith('.') && !EXCLUDED_DIRS.has(e.name))
                         .map(e => e.isDirectory() ? `${e.name}/` : e.name)
                         .join('\n');
+                    
+                    if (listing.length > 5000) {
+                        listing = listing.slice(0, 5000) + "\n... (truncated for context)";
+                    }
+
                     report.push(`📂 Listed: ${relPath}`);
                     this.chatHistory.push({ role: 'user', content: `[Result of list_files ${relPath}]\n${listing}` });
                 }
             } catch (err: any) { report.push(`❌ Listing failed: ${err.message}`); }
         }
 
+        // Action 7: Second Brain Knowledge (List/Read)
+        const listBrainRegex = /<list_brain\s*path=['"]?([^'"]*)['"]?\s*\/?>(?:<\/list_brain>)?/gi;
+        while ((match = listBrainRegex.exec(aiMessage)) !== null) {
+            const relPath = match[1].trim() || '.';
+            try {
+                const brainDir = _getBrainDir();
+                const absPath = path.join(brainDir, relPath);
+                if (fs.existsSync(absPath) && fs.statSync(absPath).isDirectory()) {
+                    const entries = fs.readdirSync(absPath, { withFileTypes: true });
+                    let listing = entries
+                        .filter(e => !e.name.startsWith('.') && !EXCLUDED_DIRS.has(e.name))
+                        .map(e => e.isDirectory() ? `${e.name}/` : e.name)
+                        .join('\n');
+                    
+                    if (listing.length > 5000) {
+                        listing = listing.slice(0, 5000) + "\n... (truncated for context)";
+                    }
+                    
+                    report.push(`🧠 Brain Listed: ${relPath}`);
+                    this.chatHistory.push({ role: 'user', content: `[Result of list_brain ${relPath}]\n${listing}` });
+                } else {
+                    report.push(`❌ Brain List failed: ${relPath} not found`);
+                }
+            } catch (err: any) { report.push(`❌ Error Listing Brain: ${err.message}`); }
+        }
+
+        const brainRegex = /<read_brain>([\s\S]*?)<\/read_brain>/gi;
+        while ((match = brainRegex.exec(aiMessage)) !== null) {
+            const fileName = match[1].trim();
+            try {
+                const brainDir = _getBrainDir();
+                const files = findBrainFiles(brainDir);
+                // Look for direct match or path match
+                const targetFile = files.find((f: string) => path.basename(f) === fileName || f.endsWith(fileName));
+                
+                if (targetFile && fs.existsSync(targetFile)) {
+                    const content = fs.readFileSync(targetFile, 'utf-8');
+                    report.push(`🧠 Brain Read: ${fileName}`);
+                    this.chatHistory.push({ role: 'user', content: `[Result of read_brain ${fileName}]\n\`\`\`\n${content}\n\`\`\`` });
+                } else {
+                    report.push(`❌ Brain Read failed: ${fileName} not found in Second Brain`);
+                }
+            } catch (err: any) { report.push(`❌ Error Reading Brain: ${err.message}`); }
+        }
+
+        // Action 8: Read URL (Simple implementation)
+        const urlRegex = /<read_url>([\s\S]*?)<\/read_url>/gi;
+        while ((match = urlRegex.exec(aiMessage)) !== null) {
+            const url = match[1].trim();
+            try {
+                const res = await fetch(url, { signal: AbortSignal.timeout(10000) });
+                const text = await res.text();
+                // Simple HTML to text-ish conversion
+                const content = text.replace(/<[^>]+>/g, ' ').replace(/\s+/g, ' ').trim();
+                const preview = content.length > 5000 ? content.slice(0, 5000) + "\n... (truncated)" : content;
+                report.push(`🌐 Read URL: ${url}`);
+                this.chatHistory.push({ role: 'user', content: `[Result of read_url ${url}]\n${preview}` });
+            } catch (err: any) { report.push(`❌ URL Read failed: ${err.message}`); }
+        }
+
         if (firstCreatedFile) {
             vscode.window.showTextDocument(vscode.Uri.file(firstCreatedFile), { preview: false });
         }
@@ -300,6 +452,8 @@ export class AgentExecutor {
             execSync(`git add .`, { cwd: brainDir });
             execSync(`git commit -m "[G1nation] Knowledge Update"`, { cwd: brainDir });
             execSync(`git push`, { cwd: brainDir });
-        } catch {}
+        } catch (err) {
+            console.error('[Agent] Sync failed:', err);
+        }
     }
 }