Fixed prompt injection warnings waiting for other actions to finish (#1053)

app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs

@@ -77,7 +77,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
         var textChunker = new TextChunker(rustService, logger);
         return
         [
-            new FileSourceIndexer(settingsManager, rustService, guardService, textChunker, logger),
+            new FileSourceIndexer(settingsManager, rustService, textChunker, logger),
             new MailboxIndexer(settingsManager, rustService, guardService, textChunker, logger),
         ];
     }
@@ -702,9 +702,10 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
     /// <param name="refreshMode">Why the run was started, for the log.</param>
     /// <param name="share">How much work the run does at most, as its kind of data source says.</param>
     /// <param name="runControl">The run, which records the embedding settings it works for.</param>
+    /// <param name="promptInjectionScope">The run as one action of the prompt injection filter.</param>
     /// <param name="token">The cancellation token.</param>
     /// <returns>The context of the run, or null when there is no run.</returns>
-    private async Task<IndexedRunContext?> PrepareIndexedRunAsync(IIndexedDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, RunShare share, DataSourceRunControl runControl, CancellationToken token)
+    private async Task<IndexedRunContext?> PrepareIndexedRunAsync(IIndexedDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, RunShare share, DataSourceRunControl runControl, PromptInjectionReportingScope promptInjectionScope, CancellationToken token)
     {
         logger.LogInformation(
             "Starting background embedding hash check for data source '{DataSourceName}' ({DataSourceId}). RefreshMode={RefreshMode}.",
@@ -831,7 +832,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
                 pass.FailedDocumentCount,
                 refreshMode);
 
-        return new IndexedRunContext(dataSource, embeddingProvider, embeddingProvider.CreateProvider(), vectorStore, indexStore, manifest, share, pass, settingsManager, this.UpsertStatus, logger);
+        return new IndexedRunContext(dataSource, embeddingProvider, embeddingProvider.CreateProvider(), vectorStore, indexStore, manifest, share, pass, promptInjectionScope, settingsManager, this.UpsertStatus, logger);
     }
 
     private async Task ProcessDataSourceAsync(IIndexedDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, DataSourceRunControl runControl, CancellationToken token)
@@ -848,7 +849,13 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
         var outcome = IndexedRunOutcome.DONE;
         try
         {
-            var context = await this.PrepareIndexedRunAsync(dataSource, refreshMode, indexer.Share, runControl, token);
+            //
+            // Everything the prompt injection filter removes from the documents of this run is
+            // reported once, when the run is over. A step over a few thousand documents which
+            // removes something in forty of them is one thing that happened to the user, not forty.
+            //
+            await using var promptInjectionScope = guardService.BeginAction();
+            var context = await this.PrepareIndexedRunAsync(dataSource, refreshMode, indexer.Share, runControl, promptInjectionScope, token);
             if (context is null)
                 return;
 

app/MindWork AI Studio/Tools/Services/Indexing/AGENTS.md

@@ -16,9 +16,9 @@ The parts:
   finds and reads the documents of one run, and `TrackChanges` / `StopTracking` notice changes on their own.
   `FileSourceIndexer` is the reference: a file system watcher per data source, fingerprints over path, size and
   write time.
-- **`IndexedRunContext`** - one prepared run: both stores, the embedding provider, the manifest and the
-  collection. `IndexDocumentAsync` embeds and stores one document; the cleanup methods remove what a failed
-  attempt left behind.
+- **`IndexedRunContext`** - one prepared run: both stores, the embedding provider, the manifest, the
+  collection, and the scope for reporting prompt injections. `IndexDocumentAsync` embeds and stores one
+  document; the cleanup methods remove what a failed attempt left behind.
 - **`EmbeddingDocument`** - one document: its key, its index row, its display name, and how to read its chunks.
 - **`DocumentRunProgress`** - counts the documents, records indexed and failed ones in the stores, publishes the
   status, and completes or pauses the run.
@@ -49,6 +49,10 @@ To add a kind of data source:
    returns: it decides whether the data source carries on in another run, or stops because its embedding
    provider failed for everything. A run which takes more than twice its share without asking once ends up as
    a warning in the log, while every other data source waits.
+7. Filter everything the indexer reads: files through `ExtractedFileText.ReadAsync`, any other content through
+   `PromptInjectionGuardService.SanitizeAsync`, both with `IndexedRunContext.PromptInjectionScope`. The service
+   opens that scope per run, so the user hears about a run once, when it is over. Begin no action of your own
+   in the indexer, or the user hears about one run several times.
 
 Rules which are easy to break:
 - **A document key is not a path.** Only files use their full path as the key. Never pass a key through the

app/MindWork AI Studio/Tools/Services/Indexing/ExtractedFileText.cs

@@ -18,13 +18,14 @@ internal static class ExtractedFileText
     /// <param name="filePath">The path of the file.</param>
     /// <param name="embeddingProvider">The embedding provider whose tokenizer counts the pieces.</param>
     /// <param name="reportAs">The source the user is told about when passages were filtered out, or null for the file itself.</param>
+    /// <param name="reportingScope">The action the file is read for, which tells the user about filtered passages together with the rest of it.</param>
     /// <param name="token">The cancellation token.</param>
     /// <returns>The text, empty when the file has none.</returns>
-    public static async Task<SegmentedText> ReadAsync(RustService rustService, string filePath, EmbeddingProvider embeddingProvider, PromptInjectionSource? reportAs, CancellationToken token)
+    public static async Task<SegmentedText> ReadAsync(RustService rustService, string filePath, EmbeddingProvider embeddingProvider, PromptInjectionSource? reportAs, PromptInjectionReportingScope reportingScope, CancellationToken token)
     {
         var segments = new List<TextSegment>();
 
-        await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, reportAs, token))
+        await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, reportAs, reportingScope, token))
         {
             var normalized = TextChunker.NormalizeSegment(segment.Content);
             if (!string.IsNullOrWhiteSpace(normalized))

app/MindWork AI Studio/Tools/Services/Indexing/FileSourceIndexer.cs

@@ -32,10 +32,9 @@ namespace AIStudio.Tools.Services.Indexing;
 /// </remarks>
 /// <param name="settingsManager">The settings, read for whether local data sources refresh on their own.</param>
 /// <param name="rustService">The runtime, which extracts the text of the files.</param>
-/// <param name="guardService">The prompt injection filter, whose findings are reported once per run.</param>
 /// <param name="textChunker">Cuts the text of a file into chunks.</param>
 /// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
-internal sealed partial class FileSourceIndexer(SettingsManager settingsManager, RustService rustService, PromptInjectionGuardService guardService, TextChunker textChunker, ILogger logger) : IIndexedSourceIndexer
+internal sealed partial class FileSourceIndexer(SettingsManager settingsManager, RustService rustService, TextChunker textChunker, ILogger logger) : IIndexedSourceIndexer
 {
     private const string OFFICE_LOCK_FILE_PREFIX = "~$";
     private const bool IMAGE_EMBEDDING_ENABLED = false;
@@ -166,14 +165,6 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
             partition.ToRead.Count);
         progress.Publish();
 
-        //
-        // Everything the runtime filters out of these files is reported once for the whole step.
-        // A step over a few thousand documents which removes something in forty of them is one
-        // thing that happened to the user, not forty. The scope ends with this method, so the
-        // report arrives when the step is over rather than in the middle of it.
-        //
-        await using var promptInjectionReportingScope = guardService.BeginAction();
-
         for (var fileIndex = 0; fileIndex < partition.ToRead.Count; fileIndex++)
         {
             token.ThrowIfCancellationRequested();
@@ -440,11 +431,11 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
             : string.Join("; ", reasons);
     }
 
-    private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
+    private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, PromptInjectionReportingScope reportingScope, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
     {
         var options = DataSourceEmbeddingService.GetChunkingOptions(dataSource, embeddingProvider);
         var strategy = TextChunker.GetStrategyForFile(filePath);
-        var content = await ExtractedFileText.ReadAsync(rustService, filePath, embeddingProvider, null, token);
+        var content = await ExtractedFileText.ReadAsync(rustService, filePath, embeddingProvider, null, reportingScope, token);
 
         await foreach (var chunk in textChunker.SplitAsync(content, strategy, options, embeddingProvider, token))
             yield return chunk;
@@ -731,7 +722,7 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
         file.FullName,
         this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTimeOffset.UtcNow),
         file.Name,
-        token => this.StreamEmbeddingChunksAsync(file.FullName, dataSource, context.EmbeddingProvider, token));
+        token => this.StreamEmbeddingChunksAsync(file.FullName, dataSource, context.EmbeddingProvider, context.PromptInjectionScope, token));
 
     private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, FileInfo file, string fingerprint, int chunkCount, DateTimeOffset embeddedAtUtc)
     {

app/MindWork AI Studio/Tools/Services/Indexing/IndexedRunContext.cs

@@ -3,6 +3,7 @@ using AIStudio.Settings;
 using AIStudio.Tools.Databases.IndexStore;
 using AIStudio.Tools.Databases.VectorStore;
 using AIStudio.Tools.PluginSystem;
+using AIStudio.Tools.Security;
 
 namespace AIStudio.Tools.Services.Indexing;
 
@@ -26,10 +27,11 @@ namespace AIStudio.Tools.Services.Indexing;
 /// <param name="manifest">What the index stores about the data source, made to match the current embedding configuration.</param>
 /// <param name="share">How much work the run does at most, as its kind of data source says.</param>
 /// <param name="pass">The pass the run is a step of, which knows what failed during the steps before.</param>
+/// <param name="promptInjectionScope">The run as one action of the prompt injection filter, which tells the user about everything it removed once the run is over.</param>
 /// <param name="settingsManager">The settings, which the provider reads when it embeds.</param>
 /// <param name="publishStatus">Hands a new status of the data source to the embedding service.</param>
 /// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
-internal sealed class IndexedRunContext(IIndexedDataSource dataSource, EmbeddingProvider embeddingProvider, IProvider provider, VectorStoreClient vectorStore, IndexStoreClient indexStore, DataSourceEmbeddingManifest manifest, RunShare share, IndexingPass pass, SettingsManager settingsManager, Action<DataSourceEmbeddingStatus> publishStatus, ILogger logger)
+internal sealed class IndexedRunContext(IIndexedDataSource dataSource, EmbeddingProvider embeddingProvider, IProvider provider, VectorStoreClient vectorStore, IndexStoreClient indexStore, DataSourceEmbeddingManifest manifest, RunShare share, IndexingPass pass, PromptInjectionReportingScope promptInjectionScope, SettingsManager settingsManager, Action<DataSourceEmbeddingStatus> publishStatus, ILogger logger)
 {
     /// <summary>
     /// After how many stored chunks the collection is optimized while a run is still going.
@@ -64,6 +66,16 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
 
     public IndexingPass Pass => pass;
 
+    /// <summary>
+    /// The run as one action of the prompt injection filter.
+    /// </summary>
+    /// <remarks>
+    /// Everything read for the run carries it down to the filter, so a run over a few thousand
+    /// documents which removes something in forty of them tells the user once, when it is over,
+    /// rather than forty times. Every kind of data source gets this from here, so none can forget it.
+    /// </remarks>
+    public PromptInjectionReportingScope PromptInjectionScope => promptInjectionScope;
+
     /// <summary>
     /// When the run first asked whether it did its share of the work, or null while it never did.
     /// </summary>

app/MindWork AI Studio/Tools/Services/Indexing/MailboxIndexer.Attachments.cs

@@ -45,7 +45,7 @@ internal sealed partial class MailboxIndexer
                 await connector.FetchAttachmentAsync(uid, part, file, token);
 
             // Filtered by the runtime while it reads, and reported as content of the mailbox, not of a file the user never saw:
-            var content = await ExtractedFileText.ReadAsync(rustService, path, context.EmbeddingProvider, PromptInjectionSource.MailContent(mailbox.Name), token);
+            var content = await ExtractedFileText.ReadAsync(rustService, path, context.EmbeddingProvider, PromptInjectionSource.MailContent(mailbox.Name), context.PromptInjectionScope, token);
             return new(part, name, MailPartTextState.EXTRACTED, content, TextChunker.GetStrategyForFile(path));
         }
         catch (OperationCanceledException) when (token.IsCancellationRequested)

app/MindWork AI Studio/Tools/Services/Indexing/MailboxIndexer.Mails.cs

@@ -148,7 +148,7 @@ internal sealed partial class MailboxIndexer
             new(text.Body, source),
             new(text.Subject, source),
             ..attachmentParts.Select(part => new PromptInjectionText(MailTextNormalization.NormalizeHeaderValue(part.FileName), source)),
-        ]);
+        ], context.PromptInjectionScope);
 
         var attachments = new List<MailAttachmentText>(attachmentParts.Count);
         for (var index = 0; index < attachmentParts.Count; index++)

app/MindWork AI Studio/Tools/Services/Indexing/MailboxIndexer.cs

@@ -244,22 +244,15 @@ internal sealed partial class MailboxIndexer(SettingsManager settingsManager, Ru
             progress.RecordUnchanged(plannedFolders.Sum(folder => folder.Plan.KeptUids.Count));
             progress.Publish();
 
-            //
-            // Whatever the prompt injection filter removes from the mails of this run is reported
-            // once, when the second pass is done, rather than once per mail.
-            //
             var encounteredKeys = new HashSet<string>(StringComparer.Ordinal);
             var workedThrough = true;
-            await using (guardService.BeginAction())
+            foreach (var folder in plannedFolders)
             {
-                foreach (var folder in plannedFolders)
-                {
-                    if (await this.SyncNewMailsAsync(context, mailbox, connector, folder, progress, encounteredKeys, token))
-                        continue;
+                if (await this.SyncNewMailsAsync(context, mailbox, connector, folder, progress, encounteredKeys, token))
+                    continue;
 
-                    workedThrough = false;
-                    break;
-                }
+                workedThrough = false;
+                break;
             }
 
             var sourceHash = BuildSourceHash(plannedFolders);

app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs

@@ -37,7 +37,8 @@ public sealed partial class RustService
     /// attachment: nothing leaves the app on that path, so there is nothing to warn about, and
     /// reporting it there would warn a second time when the file is actually sent.
     /// </param>
-    public async Task<FileExtractionResult> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false, bool reportPromptInjections = true, CancellationToken token = default)
+    /// <param name="reportingScope">The action the file is read for, or null when it is read for none.</param>
+    public async Task<FileExtractionResult> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false, bool reportPromptInjections = true, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
     {
         //
         // The runtime filters prompt injections while it streams the file. Doing it there rather
@@ -249,7 +250,7 @@ public sealed partial class RustService
         // where the content never leaves the app.
         //
         if (reportPromptInjections)
-            await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
+            await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount), reportingScope);
 
         //
         // Filtering does not change the outcome: the passages were removed and the document
@@ -263,12 +264,6 @@ public sealed partial class RustService
         };
     }
 
-    public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
-    {
-        await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, string.Empty, null, token))
-            yield return segment.Content;
-    }
-
     /// <summary>
     /// Reads the content of a file for embedding, each piece together with its token count.
     /// </summary>
@@ -279,15 +274,17 @@ public sealed partial class RustService
     /// itself. A file which only exists while it is read, e.g. the attachment of a mail, names
     /// where its content came from instead.
     /// </param>
+    /// <param name="reportingScope">The action the file is read for, or null when it is read for none.</param>
     /// <param name="token">The cancellation token.</param>
     /// <returns>The pieces of the content, in order.</returns>
     public async IAsyncEnumerable<ArbitraryFileDataSegment> StreamArbitraryFileDataWithTokenCounts(
         string path,
         EmbeddingProvider embeddingProvider,
         PromptInjectionSource? reportAs = null,
+        PromptInjectionReportingScope? reportingScope = null,
         [EnumeratorCancellation] CancellationToken token = default)
     {
-        await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, reportAs, token))
+        await foreach (var segment in this.StreamArbitraryFileDataCore(path, embeddingProvider.TokenizerPath, reportAs, reportingScope, token))
         {
             if (segment.TokenCount is { } tokenCount)
             {
@@ -318,14 +315,13 @@ public sealed partial class RustService
 
     private async IAsyncEnumerable<(string Content, int? TokenCount, int? PageNumber)> StreamArbitraryFileDataCore(
         string path,
-        bool extractImages,
-        bool includeTokenCount,
         string tokenizerPath,
         PromptInjectionSource? reportAs,
+        PromptInjectionReportingScope? reportingScope,
         [EnumeratorCancellation] CancellationToken token)
     {
         var streamId = Guid.NewGuid().ToString();
-        var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}&include_token_count={includeTokenCount}&tokenizer_path={Uri.EscapeDataString(tokenizerPath)}";
+        var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images=false&include_token_count=true&tokenizer_path={Uri.EscapeDataString(tokenizerPath)}";
         using var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
         using var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, token);
 
@@ -339,10 +335,7 @@ public sealed partial class RustService
                 path,
                 responseBody);
 
-            if (includeTokenCount)
-                throw new InvalidOperationException($"Rust could not extract and count '{path}'. HTTP {(int)response.StatusCode} ({response.ReasonPhrase}): {responseBody}");
-
-            yield break;
+            throw new InvalidOperationException($"Rust could not extract and count '{path}'. HTTP {(int)response.StatusCode} ({response.ReasonPhrase}): {responseBody}");
         }
 
         var promptInjectionFindings = new List<PromptInjectionFinding>();
@@ -372,12 +365,7 @@ public sealed partial class RustService
                 catch (JsonException)
                 {
                     if (this.TryLogSseErrorMessage(jsonContent, path))
-                    {
-                        if (includeTokenCount)
-                            throw new InvalidOperationException($"Rust could not extract and count a segment from '{path}'. See the runtime log for details.");
-
-                        continue;
-                    }
+                        throw new InvalidOperationException($"Rust could not extract and count a segment from '{path}'. See the runtime log for details.");
 
                     this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
                 }
@@ -385,7 +373,7 @@ public sealed partial class RustService
                 if (sseEvent is null)
                     continue;
 
-                var processedEvent = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages);
+                var processedEvent = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages: false);
                 if (processedEvent.Error is { } error)
                 {
                     // A notice says something about the file without failing the read, so the
@@ -461,7 +449,7 @@ public sealed partial class RustService
         // indexing run could filter documents without ever saying so.
         //
         var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
-        await guardService.ReportAsync(new(reportAs ?? PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
+        await guardService.ReportAsync(new(reportAs ?? PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount), reportingScope);
     }
 
     private bool TryLogSseErrorMessage(string jsonContent, string path)