Add MAX_STREAM_CHUNKS constant for streaming control

Improve path validation to detect parent directory traversal
Refactor VaultIndexer to use VaultLike interface
Enhance heading matching to support 1-6 level headings
Improve token stemming with length-based checks
Optimize context building for user messages
This commit is contained in:
2026-05-06 19:58:26 +02:00
parent d57ade0ba8
commit 0a0173363b
3 changed files with 43 additions and 103 deletions
+34 -25
View File
@@ -8,6 +8,7 @@ type MouseEvent = globalThis.MouseEvent;
const DEFAULT_VAULT_SEARCH_LIMIT = 3;
const MAX_MESSAGE_HISTORY = 50;
const MAX_STREAM_CHUNKS = 1000;
import {
PluginSettings,
OllamaMessage,
@@ -204,7 +205,10 @@ export class ChatView extends ItemView {
void this.inputKeyDownHandler?.(e);
};
(this.sendButton as HTMLButtonElement).addEventListener('click', this.sendButtonEventHandler);
(this.inputEl as HTMLTextAreaElement).addEventListener('keydown', this.inputKeyDownEventHandler);
(this.inputEl as HTMLTextAreaElement).addEventListener(
'keydown',
this.inputKeyDownEventHandler
);
if (this.newChatButton) {
if (!this.newChatButtonClickHandler) {
this.newChatButtonClickHandler = () => this.clearConversation();
@@ -288,9 +292,7 @@ export class ChatView extends ItemView {
// Search vault using user message as query
const entries = await this.vaultIndexer.searchVault(userMessage, DEFAULT_VAULT_SEARCH_LIMIT);
let context = entries
.map((entry) => `### ${entry.title}\n${entry.content}`)
.join('\n\n');
let context = entries.map((entry) => `### ${entry.title}\n${entry.content}`).join('\n\n');
// Cap context size to prevent prompt bloat with large vaults
const MAX_CONTEXT_LENGTH = 4000;
@@ -302,9 +304,10 @@ export class ChatView extends ItemView {
role: 'system',
content: 'You are a helpful assistant.',
};
const userContent = context ? `${context}\n\n${userMessage}` : userMessage;
const userMessageWithContext: OllamaMessage = {
role: 'user',
content: `${context}\n\n${userMessage}`,
content: userContent,
};
const messages: OllamaMessage[] = [
@@ -329,8 +332,11 @@ export class ChatView extends ItemView {
parameters: {
type: 'object' as const,
properties: {
path: { type: 'string' as const },
content: { type: 'string' as const },
path: {
type: 'string' as const,
description: "Relative path within the vault, e.g. 'Notes/todo.md'",
},
content: { type: 'string' as const, description: 'Content of the file to create' },
},
required: ['path', 'content'],
},
@@ -369,12 +375,9 @@ export class ChatView extends ItemView {
let fullResponse = '';
let toolCalls: ToolCall[] = [];
let chunkCount = 0;
const MAX_STREAM_CHUNKS = 1000;
const maxChunks = MAX_STREAM_CHUNKS;
for await (const chunk of stream) {
chunkCount++;
if (chunkCount > maxChunks) {
if (chunkCount > MAX_STREAM_CHUNKS) {
throw new Error('Response too long, stopped streaming');
}
@@ -418,22 +421,28 @@ export class ChatView extends ItemView {
}
}
// Create follow-up messages including the assistant's tool calls and results
const followUpMessages: OllamaMessage[] = [
...messages,
{ role: 'assistant' as const, content: fullResponse, tool_calls: toolCalls },
...toolResults.map((result) => ({
role: 'tool' as const,
content: JSON.stringify(result),
})),
];
// Only create follow-up when we have tool results
if (toolResults.length > 0) {
// Create follow-up messages including the assistant's tool calls and results
const followUpMessages: OllamaMessage[] = [
...messages,
{ role: 'assistant' as const, content: fullResponse, tool_calls: toolCalls },
...toolResults.map((result) => ({
role: 'tool' as const,
content: JSON.stringify(result),
})),
];
const followUp = await this.ollamaClient.chat(followUpMessages, tools);
fullResponse += followUp.content;
this.updateLastMessage(fullResponse);
const followUp = await this.ollamaClient.chat(followUpMessages, tools);
fullResponse += followUp.content;
this.updateLastMessage(fullResponse);
// Update the assistant message with the final response immutably
this.updateMessageById(assistantMessageId, { content: fullResponse, isStreaming: false });
// Update the assistant message with the final response immutably
this.updateMessageById(assistantMessageId, {
content: fullResponse,
isStreaming: false,
});
}
}
// Update last message immutably — only if no tool calls were processed
+1 -1
View File
@@ -51,7 +51,7 @@ export class ToolExecutor {
// Reject paths that traverse to parent directories
const normalized = path.replace(/^(\.\/)+/, '');
if (normalized.includes('../')) {
if (normalized.split('/').includes('..')) {
return false;
}
+8 -77
View File
@@ -1,8 +1,6 @@
// src/vault-indexer.ts
import { OllamaClient } from './ollama-client';
import { ValidationError, VaultIndexEntry } from './types';
import { ErrorHandler } from './error-handler';
import { VaultIndexEntry } from './types';
import { Logger, sanitizeFilePath } from './utils';
interface Frontmatter {
@@ -44,80 +42,11 @@ function isVaultLike(value: unknown): value is VaultLike {
}
class VaultIndexer {
private ollamaClient: OllamaClient | null = null;
private summaries: Map<string, string> = new Map();
private vault: VaultLike | null = null;
constructor(vaultOrClient: VaultLike | OllamaClient | null) {
// Support both old (OllamaClient) and new (VaultLike) interfaces
if (isVaultLike(vaultOrClient)) {
this.vault = vaultOrClient;
} else {
this.ollamaClient = vaultOrClient;
}
}
async indexVault(vaultPath: string): Promise<void> {
try {
const files = this.getMarkdownFilesInVault(vaultPath);
for (const file of files) {
const content = this.readFileContent(file.path);
const summary = await this.summarizeFile(content);
this.storeSummary(file.path, summary);
}
} catch (error: unknown) {
if (error instanceof Error) {
ErrorHandler.handleError(error, 'VaultIndexer.indexVault');
} else {
throw new ValidationError('An unexpected error occurred while indexing the vault');
}
}
}
private getMarkdownFilesInVault(vaultPath: string): { path: string }[] {
// Simulate getting markdown files from the vault
const sanitizedPath = sanitizeFilePath(vaultPath);
try {
// This is a placeholder for actual file system operations
return [{ path: `${sanitizedPath}/file1.md` }, { path: `${sanitizedPath}/file2.md` }];
} catch (error) {
throw new ValidationError('Failed to get markdown files from vault');
}
}
private readFileContent(filePath: string): string {
const sanitizedPath = sanitizeFilePath(filePath);
try {
// This is a placeholder for actual file reading operations
return `Content of ${sanitizedPath}`;
} catch (error) {
throw new ValidationError('Failed to read file content');
}
}
private async summarizeFile(content: string): Promise<string> {
if (!this.ollamaClient) {
throw new ValidationError('OllamaClient not available for summarization');
}
try {
// Use the existing chat API to summarize text
const messages = [
{
role: 'system' as const,
content: 'Summarize the following text concisely:',
},
{
role: 'user' as const,
content: content,
},
];
const response = await this.ollamaClient.chat(messages);
return response.content;
} catch (error) {
throw new ValidationError('Failed to summarize file');
}
constructor(vault: VaultLike) {
this.vault = vault;
}
private storeSummary(filePath: string, summary: string): void {
@@ -268,7 +197,7 @@ class VaultIndexer {
}
}
const headingMatches = content.match(/^# (.*?)$/gm);
const headingMatches = content.match(/^#{1,6} (.*?)$/gm);
if (headingMatches) {
headings.push(...headingMatches.map((h: string) => h.replace(/^# /, '')));
}
@@ -350,9 +279,11 @@ class VaultIndexer {
}
private stemToken(token: string): string {
// Improved stemmer that handles edge cases
if (token.length <= 3) return token; // Don't stem very short tokens
if (token.endsWith('s')) return token.slice(0, -1);
if (token.endsWith('ed')) return token.slice(0, -2);
if (token.endsWith('ing')) return token.slice(0, -3);
if (token.endsWith('ed') && token.length > 4) return token.slice(0, -2); // Don't stem 3-letter words ending in ed
if (token.endsWith('ing') && token.length > 5) return token.slice(0, -3); // Don't stem 4-letter words ending in ing
return token;
}