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

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-10-12 12:39:20 +02:00
1 parent a3c6b29313
commit 57dda9ff8b
25 files changed
+589 -146

No files matched your search

@@ -5,6 +5,7 @@ using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.RAG;
using AIStudio.Tools.Security;
using AIStudio.Tools.Services;
namespace AIStudio.Agents;
@@ -165,9 +166,10 @@ public sealed class AgentRetrievalContextValidation (ILogger<AgentRetrievalConte
/// <param name="lastUserPrompt">The last user prompt.</param>
/// <param name="chatThread">The chat thread.</param>
/// <param name="retrievalContexts">All retrieval contexts to validate.</param>
/// <param name="reportingScope">The action the contexts are validated for, or null when they are validated for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The validation results.</returns>
public async Task<IReadOnlyList<RetrievalContextValidationResult>> ValidateRetrievalContextsAsync(IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default)
public async Task<IReadOnlyList<RetrievalContextValidationResult>> ValidateRetrievalContextsAsync(IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
{
// Check if the retrieval context validation is enabled:
if (!this.SettingsManager.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation)
@@ -193,7 +195,7 @@ public sealed class AgentRetrievalContextValidation (ILogger<AgentRetrievalConte
await semaphore.WaitAsync(token);
// Start the next validation task:
validationTasks.Add(this.ValidateRetrievalContextAsync(lastUserPrompt, chatThread, retrievalContext, semaphore, token));
validationTasks.Add(this.ValidateRetrievalContextAsync(lastUserPrompt, chatThread, retrievalContext, semaphore, reportingScope, token));
}
// Wait for all validation tasks to complete:
@@ -212,9 +214,10 @@ public sealed class AgentRetrievalContextValidation (ILogger<AgentRetrievalConte
/// <param name="chatThread">The chat thread.</param>
/// <param name="retrievalContext">The retrieval context to validate.</param>
/// <param name="semaphore">The optional semaphore to limit the number of parallel validations.</param>
/// <param name="reportingScope">The action the context is validated for, or null when it is validated for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The validation result.</returns>
public async Task<RetrievalContextValidationResult> ValidateRetrievalContextAsync(IContent lastUserPrompt, ChatThread chatThread, IRetrievalContext retrievalContext, SemaphoreSlim? semaphore = null, CancellationToken token = default)
public async Task<RetrievalContextValidationResult> ValidateRetrievalContextAsync(IContent lastUserPrompt, ChatThread chatThread, IRetrievalContext retrievalContext, SemaphoreSlim? semaphore = null, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
{
try
{
@@ -252,7 +255,7 @@ public sealed class AgentRetrievalContextValidation (ILogger<AgentRetrievalConte
// 2. Prepare the retrieval context for the agent:
//
var additionalData = new Dictionary<string, string>();
var markdownRetrievalContext = await retrievalContext.AsMarkdown(token: token);
var markdownRetrievalContext = await retrievalContext.AsMarkdown(reportingScope: reportingScope, token: token);
additionalData.Add("retrievalContext", markdownRetrievalContext);
//
@@ -869,7 +869,7 @@ public partial class DocumentAnalysisAssistant : AssistantBaseCore<NoSettingsPan
continue;
}
var extraction = await this.RustService.ReadArbitraryFileData(document.FilePath, int.MaxValue);
var extraction = await this.RustService.ReadArbitraryFileData(document.FilePath, int.MaxValue, reportingScope: promptInjectionScope);
if (!extraction.HasUsableContent)
{
this.Logger.LogError("Reading the document '{FilePath}' failed and it will not be analyzed: code={ErrorCode}, message='{ErrorMessage}'.", document.FilePath, extraction.ErrorCode, extraction.ErrorMessage);
@@ -390,7 +390,7 @@ public partial class SlideAssistant : AssistantBaseCore<SettingsDialogSlideBuild
continue;
}
var extraction = await this.RustService.ReadArbitraryFileData(document.FilePath, int.MaxValue);
var extraction = await this.RustService.ReadArbitraryFileData(document.FilePath, int.MaxValue, reportingScope: promptInjectionScope);
if (!extraction.HasUsableContent)
{
this.Logger.LogError("Reading the document '{FilePath}' failed and it will not be used: code={ErrorCode}, message='{ErrorMessage}'.", document.FilePath, extraction.ErrorCode, extraction.ErrorMessage);
+7 -9
View File
@@ -400,7 +400,12 @@ public sealed class ContentText : IContent
[JsonIgnore]
public IReadOnlySet<string> RuntimeAttachmentWebAddresses { get; internal set; } = new HashSet<string>(StringComparer.Ordinal);
public async Task<string> PrepareTextContentForAI()
/// <summary>
/// Builds the text the AI gets for this message, with the content of its attached documents.
/// </summary>
/// <param name="reportingScope">The request this message is part of, which tells the user once about everything filtered out of its attachments, or null when the request reads no documents.</param>
/// <returns>The text for the AI.</returns>
public async Task<string> PrepareTextContentForAI(PromptInjectionReportingScope? reportingScope)
{
var sb = new StringBuilder();
var attachmentWebAddresses = new HashSet<string>(StringComparer.Ordinal);
@@ -447,13 +452,6 @@ public sealed class ContentText : IContent
LOGGER.LogWarning("File attachments which need Pandoc could not be processed because the Pandoc version check failed.");
}
//
// One report for the whole batch: attaching twenty documents must produce one
// dialog listing all of them, not twenty dialogs in a row.
//
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
await using var promptInjectionScope = guardService.BeginAction();
//
// The document blocks are collected separately, so we only announce attached
// files when at least one of them could actually be read. Announcing files we
@@ -475,7 +473,7 @@ public sealed class ContentText : IContent
continue;
}
var extraction = await Program.RUST_SERVICE.ReadArbitraryFileData(document.FilePath, int.MaxValue);
var extraction = await Program.RUST_SERVICE.ReadArbitraryFileData(document.FilePath, int.MaxValue, reportingScope: reportingScope);
if (!extraction.HasUsableContent)
{
LOGGER.LogError("Reading the file attachment '{FilePath}' failed and it will not be sent: code={ErrorCode}, message='{ErrorMessage}'.", document.FilePath, extraction.ErrorCode, extraction.ErrorMessage);
@@ -1,6 +1,7 @@
using AIStudio.Provider;
using AIStudio.Provider.OpenAI;
using AIStudio.Settings;
using AIStudio.Tools.Security;
namespace AIStudio.Chat;
@@ -30,6 +31,16 @@ public static class ListContentBlockExtensions
//
var canProcessImages = provider.SupportsImageInput();
//
// One report for the whole request: every message of the conversation reads its attached
// documents anew, all at the same time. Twenty documents must produce one dialog listing
// all of them, not twenty dialogs in a row. A request without documents has nothing to
// report, and needs no scope.
//
await using var reportingScope = blocks.Any(block => block.Content is ContentText text && text.FileAttachments.MayContainDocuments())
? Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>().BeginAction()
: null;
var messageTaskList = new List<Task<IMessageBase>>(blocks.Count);
foreach (var block in blocks)
{
@@ -65,7 +76,7 @@ public static class ListContentBlockExtensions
return Task.Run(async () => new TextMessage
{
Role = roleTransformer(block.Role),
Content = await text.PrepareTextContentForAI(),
Content = await text.PrepareTextContentForAI(reportingScope),
} as IMessageBase);
}
@@ -92,7 +103,7 @@ public static class ListContentBlockExtensions
Role = roleTransformer(block.Role),
Content =
[
innerTextSubContentFactory(await text.PrepareTextContentForAI()),
innerTextSubContentFactory(await text.PrepareTextContentForAI(reportingScope)),
..images,
]
} as IMessageBase;
@@ -3,4 +3,16 @@ namespace AIStudio.Chat;
public static class ListFileAttachmentExtensions
{
public static bool ContainsImages(this List<FileAttachment> attachments) => attachments.Any(attachment => attachment.IsImage);
/// <summary>
/// Whether preparing the message for the AI may read a document through the runtime.
/// </summary>
/// <remarks>
/// Every attachment but a picture counts: its type is checked again with the latest rules
/// before it is read, cf. FileAttachment.Normalize, and may turn out to be a document then.
/// Checking that here would touch every file twice, which takes long on a network drive.
/// </remarks>
/// <param name="attachments">The attachments of the message.</param>
/// <returns>True when a document may be read.</returns>
public static bool MayContainDocuments(this List<FileAttachment> attachments) => attachments.Any(attachment => !attachment.IsImage);
}
@@ -5,6 +5,7 @@ using AIStudio.Chat;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Security;
namespace AIStudio.Tools.RAG.AugmentationProcesses;
@@ -37,7 +38,14 @@ public sealed class AugmentationOne : IAugmentationProcess
}
var numTotalRetrievalContexts = retrievalContexts.Count;
//
// One report for the whole answer. With the check of which passages fit, every passage is
// filtered twice: once for that check, and once more for the model.
//
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
await using var reportingScope = guardService.BeginAction();
// Want the user to validate all retrieval contexts?
if (settings.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation && chatThread.DataSourceOptions.AutomaticValidation)
{
@@ -48,7 +56,7 @@ public sealed class AugmentationOne : IAugmentationProcess
try
{
// Let's validate all retrieval contexts:
var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, retrievalContexts, token);
var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, retrievalContexts, reportingScope, token);
if (validationResults.Count == 0)
LOGGER.LogWarning("Retrieval context validation returned no results. Continuing augmentation with all retrieved contexts.");
else
@@ -99,7 +107,7 @@ public sealed class AugmentationOne : IAugmentationProcess
sb.AppendLine();
// Let's convert all retrieval contexts to Markdown:
await retrievalContexts.AsMarkdown(sb, token);
await retrievalContexts.AsMarkdown(sb, reportingScope, token);
// Add the augmented data to the chat thread:
chatThread.AugmentedData = sb.ToString();
@@ -38,28 +38,46 @@ public static class IRetrievalContextExtensions
contextBuilder.AppendLine($"- {link}");
}
public static async Task<string> AsMarkdown(this IReadOnlyList<IRetrievalContext> retrievalContexts, StringBuilder? sb = null, CancellationToken token = default)
/// <summary>
/// Writes what the AI is told about several retrieval contexts, filtered for prompt injections.
/// </summary>
/// <param name="retrievalContexts">The contexts to describe.</param>
/// <param name="sb">The builder to write into, or null for a new one.</param>
/// <param name="reportingScope">The action the contexts are retrieved for, or null when they are retrieved for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>Everything written into the builder so far.</returns>
public static async Task<string> AsMarkdown(this IReadOnlyList<IRetrievalContext> retrievalContexts, StringBuilder? sb = null, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
{
sb ??= new StringBuilder();
var index = 0;
//
// One report for the whole retrieval run: a query may pull in dozens of contexts, and
// the user wants to know that something was filtered, not to acknowledge it per context.
//
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
await using var reportingScope = guardService.BeginAction();
await using var retrievalScope = guardService.BeginAction(within: reportingScope);
foreach(var retrievalContext in retrievalContexts)
{
index++;
await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, token);
await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, retrievalScope, token);
}
return sb.ToString();
}
public static async Task<string> AsMarkdown(this IRetrievalContext retrievalContext, StringBuilder? sb = null, int index = -1, int numTotalRetrievalContexts = -1, CancellationToken token = default)
/// <summary>
/// Writes what the AI is told about one retrieval context, filtered for prompt injections.
/// </summary>
/// <param name="retrievalContext">The context to describe.</param>
/// <param name="sb">The builder to write into, or null for a new one.</param>
/// <param name="index">The number of the context among all of them, or -1 when it stands alone.</param>
/// <param name="numTotalRetrievalContexts">How many contexts there are, or -1 when that is unknown.</param>
/// <param name="reportingScope">The action the context is retrieved for, or null when it is retrieved for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>Everything written into the builder so far.</returns>
public static async Task<string> AsMarkdown(this IRetrievalContext retrievalContext, StringBuilder? sb = null, int index = -1, int numTotalRetrievalContexts = -1, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
{
sb ??= new StringBuilder();
var contextBuilder = new StringBuilder();
@@ -141,7 +159,7 @@ public static class IRetrievalContextExtensions
//
async Task FilterWhatWeHaveSoFar()
{
var sanitized = await guardService.SanitizeAsync(contextBuilder.ToString(), source);
var sanitized = await guardService.SanitizeAsync(contextBuilder.ToString(), source, reportingScope);
contextBuilder.Clear();
contextBuilder.Append(sanitized);
}
@@ -11,7 +11,10 @@ namespace AIStudio.Tools.Security;
/// The detection itself lives in the Rust runtime. File content is filtered while the runtime
/// streams it, so it never passes through here; what this service adds is the path for content
/// the runtime does not read itself — web pages and retrieval contexts — and the reporting the
/// user sees.
/// user sees.<br/><br/>
/// The reporting works per user action, cf. BeginAction: each action tells the user once, when
/// it ends, about what was filtered for it. One service serves the whole app, so the action has
/// to travel with the content instead of being kept here.
/// </remarks>
public sealed class PromptInjectionGuardService(
RustService rustService,
@@ -27,8 +30,6 @@ public sealed class PromptInjectionGuardService(
private readonly ILogger detectionLogger = loggerFactory.CreateLogger(DETECTION_LOG_CATEGORY);
private readonly Lock reportLock = new();
private readonly List<PromptInjectionScanResult> pendingResults = [];
private int openActions;
/// <summary>
/// Filters prompt injections out of a text the runtime did not read itself, such as a web
@@ -42,8 +43,9 @@ public sealed class PromptInjectionGuardService(
/// </remarks>
/// <param name="text">The content to filter.</param>
/// <param name="source">Where the content came from, for the report shown to the user.</param>
/// <param name="reportingScope">The action the content belongs to, or null when it belongs to none.</param>
/// <returns>The content with any suspicious passages removed.</returns>
public async Task<string> SanitizeAsync(string text, PromptInjectionSource source)
public async Task<string> SanitizeAsync(string text, PromptInjectionSource source, PromptInjectionReportingScope? reportingScope = null)
{
if (string.IsNullOrWhiteSpace(text))
return text;
@@ -59,7 +61,7 @@ public sealed class PromptInjectionGuardService(
}
if (response.RedactedCount > 0)
await this.ReportAsync(new(source, response.Findings, response.RedactedCount));
await this.ReportAsync(new(source, response.Findings, response.RedactedCount), reportingScope);
return response.SanitizedText;
}
@@ -75,8 +77,9 @@ public sealed class PromptInjectionGuardService(
/// the check cannot run, every text is passed through unchanged.
/// </remarks>
/// <param name="texts">The contents to filter, each with its source.</param>
/// <param name="reportingScope">The action the contents belong to, or null when they belong to none.</param>
/// <returns>The contents with any suspicious passages removed, in the order they came in.</returns>
public async Task<IReadOnlyList<string>> SanitizeAsync(IReadOnlyList<PromptInjectionText> texts)
public async Task<IReadOnlyList<string>> SanitizeAsync(IReadOnlyList<PromptInjectionText> texts, PromptInjectionReportingScope? reportingScope = null)
{
if (texts.Count is 0)
return [];
@@ -113,7 +116,9 @@ public sealed class PromptInjectionGuardService(
//
// Findings are collected per source, not per text: a page whose content and title were
// both filtered is one thing that happened to the user, not two.
// both filtered is one thing that happened to the user, not two. Doing so here rather than
// in the report keeps both counts: two texts with the same findings would look like one
// text filtered twice there.
//
var findingsBySource = new Dictionary<PromptInjectionSource, (List<PromptInjectionFinding> Findings, int RedactedCount)>();
for (var responseIndex = 0; responseIndex < indicesToScan.Count; responseIndex++)
@@ -139,9 +144,9 @@ public sealed class PromptInjectionGuardService(
// One scope around all sources, so a search across five pages reports once instead of
// five times:
//
await using var reportingScope = this.BeginAction();
await using var batchScope = this.BeginAction(within: reportingScope);
foreach (var (source, aggregate) in findingsBySource)
await this.ReportAsync(new(source, aggregate.Findings, aggregate.RedactedCount));
await this.ReportAsync(new(source, aggregate.Findings, aggregate.RedactedCount), batchScope);
return sanitizedTexts;
}
@@ -150,25 +155,31 @@ public sealed class PromptInjectionGuardService(
/// Records what was filtered out of one piece of content and tells the user about it.
/// </summary>
/// <remarks>
/// Within a BeginAction scope the result is collected and reported together
/// with the rest of that action. Outside of one it is reported immediately: a result that
/// simply waited for the next scope would either never reach the user, or reach them as
/// part of an unrelated action later on.
/// Within an action, the result is collected and reported together with the rest of that
/// action once it ends. Without one it is reported right away, and so is a result whose
/// action has ended already: waiting for the next action would either never reach the user,
/// or reach them as part of an unrelated action later on.<br/><br/>
/// Only the action passed in counts, never what else is running. A run over a large data
/// source takes minutes, and a message sent meanwhile must not wait for it, nor end up in the
/// same report.
/// </remarks>
public async Task ReportAsync(PromptInjectionScanResult result)
/// <param name="result">What was filtered out of the content.</param>
/// <param name="reportingScope">The action the content belongs to, or null when it belongs to none.</param>
public async Task ReportAsync(PromptInjectionScanResult result, PromptInjectionReportingScope? reportingScope = null)
{
if (!result.WasFiltered)
return;
bool reportNow;
lock (this.reportLock)
{
this.pendingResults.Add(result);
reportNow = this.openActions is 0;
if (FindCollectingScope(reportingScope) is { } collectingScope)
{
collectingScope.PendingResults.Add(result);
return;
}
}
if (reportNow)
await this.ReportPendingAsync();
await this.ReportNowAsync([result]);
}
/// <summary>
@@ -177,44 +188,104 @@ public sealed class PromptInjectionGuardService(
/// </summary>
/// <remarks>
/// Results are collected until the action finishes, so the user gets one report about
/// twenty documents instead of twenty reports. Actions may nest: only the outermost one
/// reports.
/// twenty documents instead of twenty reports. An action begun within another one is part
/// of it, and only the outermost one reports. Unrelated actions report on their own, each
/// as soon as it ends, however much they overlap.<br/><br/>
/// The scope has to travel with the content: hand it to wherever that content is read or
/// filtered, such as ReportAsync, SanitizeAsync, or RustService.ReadArbitraryFileData.
/// </remarks>
/// <param name="within">The action this one is part of, or null when it is an action of its own.</param>
/// <returns>A scope that reports what was filtered once it is disposed.</returns>
public ReportingScope BeginAction()
{
lock (this.reportLock)
this.openActions++;
public PromptInjectionReportingScope BeginAction(PromptInjectionReportingScope? within = null) => new(this, within);
return new(this);
}
private async Task EndActionAsync()
{
lock (this.reportLock)
{
this.openActions--;
// An inner scope reports nothing: the action the user started is still running.
if (this.openActions > 0)
return;
}
await this.ReportPendingAsync();
}
private async Task ReportPendingAsync()
/// <summary>
/// Ends an action and reports what was collected for it.
/// </summary>
/// <remarks>
/// A scope within an action which still runs has collected nothing, since its results went to
/// the outermost scope. It collects only once the outer one has ended before it, e.g., for work
/// which outlived the action that started it. Ending a scope a second time does nothing.
/// </remarks>
/// <param name="scope">The scope which ends.</param>
internal async Task EndActionAsync(PromptInjectionReportingScope scope)
{
List<PromptInjectionScanResult> results;
lock (this.reportLock)
{
if (this.pendingResults.Count is 0)
if (!scope.IsOpen)
return;
results = [..this.pendingResults];
this.pendingResults.Clear();
scope.IsOpen = false;
results = [..scope.PendingResults];
scope.PendingResults.Clear();
}
if (results.Count > 0)
await this.ReportNowAsync(results);
}
/// <summary>
/// Finds the scope which collects a result of the given one: the outermost scope of its action
/// which is still running.
/// </summary>
/// <remarks>
/// Call under the report lock only.
/// </remarks>
/// <param name="scope">The scope the result was reported for, or null.</param>
/// <returns>The collecting scope, or null when the result has to be reported right away.</returns>
private static PromptInjectionReportingScope? FindCollectingScope(PromptInjectionReportingScope? scope)
{
PromptInjectionReportingScope? collectingScope = null;
for (; scope is not null; scope = scope.Parent)
{
if (scope.IsOpen)
collectingScope = scope;
}
return collectingScope;
}
/// <summary>
/// Combines the results of one report into one result per source, in the order the sources came in.
/// </summary>
/// <remarks>
/// An action may filter the same content twice: the classic data source search filters a passage when
/// it checks whether the passage fits the question, and once more when it hands the passage to the
/// model. Both times, the result is the same, so a result equal to one already seen for its source is
/// dropped; the user hears about the passage once. The findings carry no position, though, so two
/// passages of one file with exactly the same findings count as one as well. That mostly happens
/// where two neighboring chunks overlap, which is the same passage indeed.<br/><br/>
/// Other results of one source, such as two passages of one file or the mails of one mailbox,
/// become a single result with all their findings and the sum of the removed passages. The report
/// counts sources, and one file is one source, however many of its passages were filtered.
/// </remarks>
/// <param name="results">The results of one report.</param>
/// <returns>One result per source.</returns>
private static IReadOnlyList<PromptInjectionScanResult> CombineBySource(IReadOnlyList<PromptInjectionScanResult> results)
{
if (results.Count < 2)
return results;
return results.GroupBy(result => result.Source).Select(CombineOneSource).ToList();
static PromptInjectionScanResult CombineOneSource(IGrouping<PromptInjectionSource, PromptInjectionScanResult> resultsOfSource)
{
List<PromptInjectionScanResult> distinctResults = [];
foreach (var result in resultsOfSource)
{
if (!distinctResults.Any(known => known.RedactedCount == result.RedactedCount && known.Findings.SequenceEqual(result.Findings)))
distinctResults.Add(result);
}
return distinctResults.Count is 1
? distinctResults[0]
: new(resultsOfSource.Key, distinctResults.SelectMany(result => result.Findings).ToList(), distinctResults.Sum(result => result.RedactedCount));
}
}
private async Task ReportNowAsync(IReadOnlyList<PromptInjectionScanResult> reportedResults)
{
var results = CombineBySource(reportedResults);
var totalCount = results.Sum(result => result.RedactedCount);
this.detectionLogger.LogWarning(
"Detected and removed {PassageCount} potentially dangerous passage(s) in {SourceCount} content source(s).",
@@ -230,12 +301,4 @@ public sealed class PromptInjectionGuardService(
if (settingsManager.ConfigurationData.App.ShowPromptInjectionAlert)
await MessageBus.INSTANCE.SendMessage<PromptInjectionAlertMessage>(null, Event.SHOW_PROMPT_INJECTION_ALERT, new(results));
}
/// <summary>
/// Reports everything filtered during one user action when it goes out of scope.
/// </summary>
public sealed class ReportingScope(PromptInjectionGuardService guardService) : IAsyncDisposable
{
public async ValueTask DisposeAsync() => await guardService.EndActionAsync();
}
}
@@ -0,0 +1,36 @@
namespace AIStudio.Tools.Security;
/// <summary>
/// One user action, such as sending a message, whose filtered content is reported to the user together.
/// </summary>
/// <remarks>
/// Created through PromptInjectionGuardService.BeginAction. Content which belongs to the action carries
/// the scope down to where it is read or filtered; disposing the scope ends the action.
/// </remarks>
/// <param name="guardService">The service which reports what was filtered.</param>
/// <param name="parent">The action this scope is part of, or null when it begins an action of its own.</param>
public sealed class PromptInjectionReportingScope(PromptInjectionGuardService guardService, PromptInjectionReportingScope? parent) : IAsyncDisposable
{
/// <summary>
/// The action this scope is part of, or null when it begins an action of its own.
/// </summary>
internal PromptInjectionReportingScope? Parent => parent;
/// <summary>
/// What was filtered while this scope was the outermost one of its action still running.
/// </summary>
/// <remarks>
/// Read and written under the lock of the guard service only.
/// </remarks>
internal List<PromptInjectionScanResult> PendingResults { get; } = [];
/// <summary>
/// Gets or sets a value indicating whether the action is still running.
/// </summary>
/// <remarks>
/// Read and written under the lock of the guard service only.
/// </remarks>
internal bool IsOpen { get; set; } = true;
public async ValueTask DisposeAsync() => await guardService.EndActionAsync(this);
}
@@ -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;
@@ -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
@@ -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))
@@ -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)
{
@@ -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>
@@ -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)
@@ -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++)
@@ -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);
@@ -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)
@@ -303,7 +303,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
// each, and so on. Otherwise, the data source offered first would take the budget, and the
// others would get what it left over.
//
await using (guardService.BeginAction())
await using (var reportingScope = guardService.BeginAction())
{
var mostPassages = pages.Select(page => page.Contexts.Count).DefaultIfEmpty(0).Max();
for (var rank = 0; rank < mostPassages; rank++)
@@ -314,7 +314,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
continue;
var retrievalContext = pages[index].Contexts[rank];
var passage = await retrievalContext.AsMarkdown(index: passageCount + 1, token: token);
var passage = await retrievalContext.AsMarkdown(index: passageCount + 1, reportingScope: reportingScope, token: token);
// A passage too long for what is left makes room for shorter ones after it:
if (textContent.Length + passage.Length > MAX_RESULT_CHARACTERS)
@@ -44,6 +44,8 @@
- Fixed AI Studio sometimes warning about unsaved changes when you left a chat in which you had only selected a chat template, without writing anything yourself.
- Fixed regenerating an answer clearing the message you had started to write.
- Fixed the warning about unsaved changes in a chat appearing only by chance when you save your chats manually or have turned off workspaces. AI Studio now reliably asks before you open another chat, create a new one, or leave the chat page, so you no longer lose your changes by accident.
- Fixed the security notice about removed instructions (prompt injection) appearing only minutes later while AI Studio updated a data source in the background. The notice also mixed your own attachments with files of the data source. Now each notice appears as soon as AI Studio has read the content, for example, the attachments of your message, and it names only that content.
- Fixed the security notice about removed instructions (prompt injection) appearing several times for one answer when AI Studio checks which passages of your data sources fit your question: once for every passage it checked, and once more for all of them together. You now get one notice per answer, and a file counts once, however many of its passages were affected.
- Updated the code contributions on the supporters page, which now thank everyone who has contributed code to AI Studio so far.
- Upgraded several libraries to improve security.
- Upgraded to Rust v1.99.0