MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-09-27 03:33:37 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Report prompt injection findings from the embedding pipeline

Browse Source
This commit is contained in:
Thorsten Sommer 2026-09-06 15:26:16 +02:00
parent 1bf11c0ca7
commit 9875fbc646
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
2 changed files with 41 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
View File

@ -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();

30
app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs
View File

@ -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)