Added a prompt injection detection (#857)

Co-authored-by: Thorsten Sommer <SommerEngineering@users.noreply.github.com>

app/MindWork AI Studio/Tools/Security/PromptInjectionAlertMessage.cs

@@ -0,0 +1,17 @@
+namespace AIStudio.Tools.Security;
+
+/// <summary>
+/// Asks the UI to tell the user what was filtered out of the content they just used.
+/// </summary>
+/// <remarks>
+/// Carries every result of one user action rather than a single one. Attaching twenty
+/// documents at once must produce one dialog listing all of them, not twenty dialogs.
+/// </remarks>
+/// <param name="Results">What was filtered, per piece of content.</param>
+public sealed record PromptInjectionAlertMessage(IReadOnlyList<PromptInjectionScanResult> Results)
+{
+    /// <summary>
+    /// Gets the total number of filtered passages across all content.
+    /// </summary>
+    public int TotalRedactedCount => this.Results.Sum(result => result.RedactedCount);
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionFinding.cs

@@ -0,0 +1,31 @@
+using System.Text.Json.Serialization;
+
+namespace AIStudio.Tools.Security;
+
+/// <summary>
+/// One passage the runtime identified as a prompt-injection attempt and filtered out.
+/// </summary>
+/// <remarks>
+/// The property names are spelled out because the content stream is deserialized without a
+/// naming policy, so the names have to match what the runtime sends verbatim.
+/// </remarks>
+public sealed record PromptInjectionFinding
+{
+    /// <summary>
+    /// Which rule matched, e.g. "instruction_override".
+    /// </summary>
+    [JsonPropertyName("rule_id")]
+    public string RuleId { get; init; } = string.Empty;
+
+    /// <summary>
+    /// The rule's family, e.g. "exfiltration".
+    /// </summary>
+    [JsonPropertyName("category")]
+    public PromptInjectionFindingCategory Category { get; init; } = PromptInjectionFindingCategory.UNKNOWN;
+
+    /// <summary>
+    /// The passage as it appeared in the content, so the user can see what was removed.
+    /// </summary>
+    [JsonPropertyName("snippet")]
+    public string Snippet { get; init; } = string.Empty;
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionFindingCategory.cs

@@ -0,0 +1,19 @@
+using System.Text.Json.Serialization;
+
+namespace AIStudio.Tools.Security;
+
+[JsonConverter(typeof(PromptInjectionFindingCategoryJsonConverter))]
+public enum PromptInjectionFindingCategory
+{
+    UNKNOWN = 0,
+    OVERRIDE,
+    ROLE_OVERRIDE,
+    EXFILTRATION,
+    JAILBREAK,
+    AGENT_MANIPULATION,
+    DELIMITER_EVASION,
+    MARKUP_EVASION,
+    ENCODING_EVASION,
+    PERSISTENCE,
+    EVASION,
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionFindingCategoryExtensions.cs

@@ -0,0 +1,23 @@
+using AIStudio.Tools.PluginSystem;
+
+namespace AIStudio.Tools.Security;
+
+public static class PromptInjectionFindingCategoryExtensions
+{
+    private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(PromptInjectionFindingCategoryExtensions).Namespace, nameof(PromptInjectionFindingCategoryExtensions));
+
+    public static string GetDisplayName(this PromptInjectionFindingCategory category) => category switch
+    {
+        PromptInjectionFindingCategory.OVERRIDE => TB("Attempt to override instructions"),
+        PromptInjectionFindingCategory.ROLE_OVERRIDE => TB("Attempt to change the AI's role"),
+        PromptInjectionFindingCategory.EXFILTRATION => TB("Attempt to expose protected data"),
+        PromptInjectionFindingCategory.JAILBREAK => TB("Attempt to bypass safeguards"),
+        PromptInjectionFindingCategory.AGENT_MANIPULATION => TB("Attempt to manipulate an agent"),
+        PromptInjectionFindingCategory.DELIMITER_EVASION => TB("Hidden instructions using delimiters"),
+        PromptInjectionFindingCategory.MARKUP_EVASION => TB("Hidden instructions using markup"),
+        PromptInjectionFindingCategory.ENCODING_EVASION => TB("Hidden instructions using encoding"),
+        PromptInjectionFindingCategory.PERSISTENCE => TB("Persistent or delayed instruction"),
+        PromptInjectionFindingCategory.EVASION => TB("Obfuscated instruction"),
+        _ => TB("Unknown"),
+    };
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionFindingCategoryJsonConverter.cs

@@ -0,0 +1,51 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace AIStudio.Tools.Security;
+
+/// <summary>
+/// Reads the finding category in the snake_case spelling the Rust runtime sends.
+/// </summary>
+/// <remarks>
+/// The converter sits on the enum itself because neither path that reads a finding passes
+/// JsonSerializerOptions: the sanitize response is read by RustService.SanitizePromptInjections
+/// and the content stream by RustService.ReadFileContent. The shared RustEnumConverter therefore
+/// never applies here, and without a converter on the type only numbers would be accepted.
+///
+/// An unrecognized category falls back to UNKNOWN instead of throwing. Throwing would cost more
+/// than the label: it fails the whole response, and the guard service then passes the content
+/// through unfiltered rather than losing a single name.
+/// </remarks>
+public sealed class PromptInjectionFindingCategoryJsonConverter : JsonConverter<PromptInjectionFindingCategory>
+{
+    private static readonly ILogger<PromptInjectionFindingCategoryJsonConverter> LOG = Program.LOGGER_FACTORY.CreateLogger<PromptInjectionFindingCategoryJsonConverter>();
+
+    public override PromptInjectionFindingCategory Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+    {
+        if (reader.TokenType is not JsonTokenType.String)
+        {
+            LOG.LogWarning("Cannot read a prompt injection finding category from a '{TokenType}' token. Using UNKNOWN.", reader.TokenType);
+            return PromptInjectionFindingCategory.UNKNOWN;
+        }
+
+        var text = reader.GetString();
+        if (string.IsNullOrWhiteSpace(text))
+        {
+            LOG.LogWarning("Read an empty prompt injection finding category. Using UNKNOWN.");
+            return PromptInjectionFindingCategory.UNKNOWN;
+        }
+
+        //
+        // The enum members are the wire value in upper case, so upper-casing replaces a naming
+        // policy. Values starting with a digit or sign are rejected up front, because Enum.TryParse
+        // would otherwise accept "0" or "-1" as a category:
+        //
+        if (!char.IsAsciiDigit(text[0]) && text[0] is not ('-' or '+') && Enum.TryParse<PromptInjectionFindingCategory>(text.ToUpperInvariant(), out var category))
+            return category;
+
+        LOG.LogWarning("The runtime reported the unknown prompt injection finding category '{Category}'. Using UNKNOWN.", text);
+        return PromptInjectionFindingCategory.UNKNOWN;
+    }
+
+    public override void Write(Utf8JsonWriter writer, PromptInjectionFindingCategory value, JsonSerializerOptions options) => writer.WriteStringValue(value.ToString().ToLowerInvariant());
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionGuardService.cs

@@ -0,0 +1,159 @@
+using AIStudio.Settings;
+using AIStudio.Tools.PluginSystem;
+using AIStudio.Tools.Services;
+
+namespace AIStudio.Tools.Security;
+
+/// <summary>
+/// Filters prompt injections out of external content before it reaches a model.
+/// </summary>
+/// <remarks>
+/// 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.
+/// </remarks>
+public sealed class PromptInjectionGuardService(
+    RustService rustService,
+    SettingsManager settingsManager,
+    ILogger<PromptInjectionGuardService> logger,
+    ILoggerFactory loggerFactory)
+{
+    public const string WIKI_URL = "https://en.wikipedia.org/wiki/Prompt_engineering#Prompt_injection";
+
+    private const string DETECTION_LOG_CATEGORY = "PromptInjectionProtection";
+
+    private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(PromptInjectionGuardService).Namespace, nameof(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
+    /// page or a retrieval context.
+    /// </summary>
+    /// <remarks>
+    /// Returns usable text in every case. When the runtime cannot be reached, the text is passed
+    /// through unchanged: refusing the user's content because a check could not run would cost
+    /// them their work over a check that is best-effort anyway. The failure is logged and shown,
+    /// so it does not pass silently.
+    /// </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>
+    /// <returns>The content with any suspicious passages removed.</returns>
+    public async Task<string> SanitizeAsync(string text, PromptInjectionSource source)
+    {
+        if (string.IsNullOrWhiteSpace(text))
+            return text;
+
+        if (await rustService.SanitizePromptInjections(text) is not { } response)
+        {
+            logger.LogError("Could not check {SourceKind} '{SourceLabel}' for prompt injections. The content is used unchanged.", source.Kind, source.Label);
+            await MessageBus.INSTANCE.SendWarning(new(
+                Icons.Material.Filled.GppMaybe,
+                string.Format(TB("AI Studio could not check '{0}' for prompt injections. The content is used as it is."), source.NotificationLabel)));
+
+            return text;
+        }
+
+        if (response.RedactedCount > 0)
+            await this.ReportAsync(new(source, response.Findings, response.RedactedCount));
+
+        return response.SanitizedText;
+    }
+
+    /// <summary>
+    /// 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.
+    /// </remarks>
+    public async Task ReportAsync(PromptInjectionScanResult result)
+    {
+        if (!result.WasFiltered)
+            return;
+
+        bool reportNow;
+        lock (this.reportLock)
+        {
+            this.pendingResults.Add(result);
+            reportNow = this.openActions is 0;
+        }
+
+        if (reportNow)
+            await this.ReportPendingAsync();
+    }
+
+    /// <summary>
+    /// Marks the start of one user action, such as attaching a batch of files or sending a
+    /// message.
+    /// </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.
+    /// </remarks>
+    /// <returns>A scope that reports what was filtered once it is disposed.</returns>
+    public ReportingScope BeginAction()
+    {
+        lock (this.reportLock)
+            this.openActions++;
+
+        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()
+    {
+        List<PromptInjectionScanResult> results;
+        lock (this.reportLock)
+        {
+            if (this.pendingResults.Count is 0)
+                return;
+
+            results = [..this.pendingResults];
+            this.pendingResults.Clear();
+        }
+
+        var totalCount = results.Sum(result => result.RedactedCount);
+        this.detectionLogger.LogWarning(
+            "Detected and removed {PassageCount} potentially dangerous passage(s) in {SourceCount} content source(s).",
+            totalCount,
+            results.Count);
+
+        await MessageBus.INSTANCE.SendWarning(new(
+            Icons.Material.Filled.GppMaybe,
+            results.Count is 1
+                ? string.Format(TB("AI Studio removed suspicious instructions from '{0}' before using it."), results[0].Source.NotificationLabel)
+                : string.Format(TB("AI Studio removed suspicious instructions from {0} sources before using them."), results.Count)));
+
+        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();
+    }
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionScanResult.cs

@@ -0,0 +1,19 @@
+namespace AIStudio.Tools.Security;
+
+/// <summary>
+/// What the runtime filtered out of one piece of external content.
+/// </summary>
+/// <param name="Source">Where the content came from, so the user can tell which file or page it was.</param>
+/// <param name="Findings">The passages that were removed. Capped by the runtime.</param>
+/// <param name="RedactedCount">How many passages were removed in total, which may exceed the number of findings.</param>
+public sealed record PromptInjectionScanResult(PromptInjectionSource Source, IReadOnlyList<PromptInjectionFinding> Findings, int RedactedCount)
+{
+    /// <summary>
+    /// Gets a value indicating whether anything was filtered out of this content.
+    /// </summary>
+    /// <remarks>
+    /// The content itself stays usable either way: passages are removed, the content around
+    /// them is not rejected.
+    /// </remarks>
+    public bool WasFiltered => this.RedactedCount > 0;
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs

@@ -0,0 +1,16 @@
+namespace AIStudio.Tools.Security;
+
+public readonly record struct PromptInjectionSource(PromptInjectionSourceKind Kind, string Label)
+{
+    public string NotificationLabel => this.Kind is PromptInjectionSourceKind.FILE_CONTENT or PromptInjectionSourceKind.CHAT_ATTACHMENT
+        ? Path.GetFileName(this.Label)
+        : this.Label;
+
+    public static PromptInjectionSource WebContent(string url) => new(PromptInjectionSourceKind.WEB_CONTENT, url);
+
+    public static PromptInjectionSource FileContent(string filePath) => new(PromptInjectionSourceKind.FILE_CONTENT, filePath);
+
+    public static PromptInjectionSource ChatAttachment(string filePath) => new(PromptInjectionSourceKind.CHAT_ATTACHMENT, filePath);
+
+    public static PromptInjectionSource RetrievalContext(string dataSourceName, string path) => new(PromptInjectionSourceKind.RETRIEVAL_CONTEXT, $"{dataSourceName}: {path}");
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKind.cs

@@ -0,0 +1,10 @@
+namespace AIStudio.Tools.Security;
+
+public enum PromptInjectionSourceKind
+{
+    UNKNOWN = 0,
+    WEB_CONTENT,
+    FILE_CONTENT,
+    CHAT_ATTACHMENT,
+    RETRIEVAL_CONTEXT,
+}

app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKindExtensions.cs

@@ -0,0 +1,17 @@
+using AIStudio.Tools.PluginSystem;
+
+namespace AIStudio.Tools.Security;
+
+public static class PromptInjectionSourceKindExtensions
+{
+    private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(PromptInjectionSourceKindExtensions).Namespace, nameof(PromptInjectionSourceKindExtensions));
+
+    public static string GetDisplayName(this PromptInjectionSourceKind kind) => kind switch
+    {
+        PromptInjectionSourceKind.WEB_CONTENT => TB("Web content"),
+        PromptInjectionSourceKind.FILE_CONTENT => TB("File content"),
+        PromptInjectionSourceKind.CHAT_ATTACHMENT => TB("Chat attachment"),
+        PromptInjectionSourceKind.RETRIEVAL_CONTEXT => TB("Retrieved context"),
+        _ => TB("Unknown"),
+    };
+}