From 9875fbc64614df462570b52fcdb28689e47691f3 Mon Sep 17 00:00:00 2001
From: Thorsten Sommer <SommerEngineering@users.noreply.github.com>
Date: Sun, 6 Sep 2026 15:26:16 +0200
Subject: [PATCH] Report prompt injection findings from the embedding pipeline

---
 .../Services/DataSourceEmbeddingService.cs    | 15 +++++++---
 .../Tools/Services/RustService.Retrieval.cs   | 30 +++++++++++++++++++
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
index e5f92458..1250313a 100644
--- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs	
+++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs	
@@ -9,13 +9,12 @@ using AIStudio.Tools.Databases;
 using AIStudio.Tools.Databases.IndexStore;
 using AIStudio.Tools.Databases.VectorStore;
 using AIStudio.Tools.PluginSystem;
+using AIStudio.Tools.Security;
 
 namespace AIStudio.Tools.Services;
 
-public sealed partial class DataSourceEmbeddingService(
-    SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider,
-    ILogger<DataSourceEmbeddingService> logger)
-    : BackgroundService
+public sealed partial class DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider,
+    PromptInjectionGuardService guardService, ILogger<DataSourceEmbeddingService> logger) : BackgroundService
 {
     private const int VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD = 100_000;
 
@@ -568,6 +567,14 @@ public sealed partial class DataSourceEmbeddingService(
         var lastError = inputFiles.LastError;
         var failureDetails = inputFiles.Failures.ToList();
 
+        //
+        // Everything the runtime filters out of these files is reported once for the whole data
+        // source. A run 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 run is finished rather than in the middle of it.
+        //
+        await using var promptInjectionReportingScope = guardService.BeginAction();
+
         foreach (var file in indexedFiles)
         {
             token.ThrowIfCancellationRequested();
diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs
index 63dea478..2cdc4b9b 100644
--- a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs	
+++ b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs	
@@ -302,6 +302,9 @@ public sealed partial class RustService
             yield break;
         }
 
+        var promptInjectionFindings = new List<PromptInjectionFinding>();
+        var promptInjectionRedactedCount = 0;
+
         string? finalContentChunk;
         try
         {
@@ -372,6 +375,21 @@ public sealed partial class RustService
                     throw new InvalidOperationException($"Rust could not extract '{path}': {error.Message}");
                 }
 
+                if (processedEvent.PromptInjection is { } promptInjection)
+                {
+                    //
+                    // Not a failure: the passages were removed and the document around them is
+                    // intact, so what remains still belongs into the index. It only has to reach
+                    // the user, because from here on the indexed document is no longer the one
+                    // sitting on their disk.
+                    //
+                    promptInjectionRedactedCount += promptInjection.RedactedCount;
+                    if (promptInjection.Findings is { } findings)
+                        promptInjectionFindings.AddRange(findings);
+
+                    continue;
+                }
+
                 if (!string.IsNullOrWhiteSpace(processedEvent.Content))
                     yield return (processedEvent.Content, sseEvent.TokenCount);
             }
@@ -383,6 +401,18 @@ public sealed partial class RustService
 
         if (!string.IsNullOrWhiteSpace(finalContentChunk))
             yield return (finalContentChunk, null);
+
+        if (promptInjectionRedactedCount is 0)
+            yield break;
+
+        //
+        // Reported from here for the same reason as in ReadArbitraryFileData above: these two
+        // methods together are every way of reading a file, so they are the only two places
+        // where no caller can forget the report. Here it was missing, which is why a whole
+        // indexing run could filter documents without ever saying so.
+        //
+        var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
+        await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
     }
 
     private bool TryLogSseErrorMessage(string jsonContent, string path)