diff --git a/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs b/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs index ee2437d9..ee409fa7 100644 --- a/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs +++ b/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs @@ -4,6 +4,7 @@ using AIStudio.Chat; using AIStudio.Provider; using AIStudio.Settings; using AIStudio.Tools.RAG; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; namespace AIStudio.Agents; @@ -237,7 +238,20 @@ public sealed class AgentRetrievalContextValidation (ILogger(); - var markdownRetrievalContext = await retrievalContext.AsMarkdown(token: token); + string markdownRetrievalContext; + try + { + markdownRetrievalContext = await retrievalContext.AsMarkdown(token: token); + } + catch (PromptInjectionBlockedException exception) + { + logger.LogWarning( + exception, + "Blocked retrieval context '{DataSourceName}' at '{Path}' before validation.", + retrievalContext.DataSourceName, + retrievalContext.Path); + return new(false, "The retrieval context was blocked due to suspected prompt injection.", 1.0f, retrievalContext); + } additionalData.Add("retrievalContext", markdownRetrievalContext); // diff --git a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs index 395f8055..6e911a90 100644 --- a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs +++ b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs @@ -5,6 +5,7 @@ using AIStudio.Dialogs.Settings; using AIStudio.Tools.AIJobs; using AIStudio.Tools.AssistantSessions; using AIStudio.Tools.Media; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; using Microsoft.AspNetCore.Components; @@ -273,6 +274,12 @@ public abstract partial class AssistantBase : AssistantLowerBase wher this.Logger.LogError(e, "The provider request failed for assistant '{AssistantTitle}'. Status={StatusCode}, Reason='{ReasonPhrase}', Body='{ResponseBody}'", this.Title, e.StatusCode, e.ReasonPhrase, e.ResponseBody); await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, e.UserMessage)); } + catch (PromptInjectionBlockedException e) + { + sessionStatus = AssistantSessionStatus.FAILED; + errorMessage = e.Message; + this.Logger.LogWarning(e, "Blocked prompt-injection content during assistant session '{AssistantTitle}'.", this.Title); + } catch (Exception e) { sessionStatus = AssistantSessionStatus.FAILED; @@ -461,6 +468,18 @@ public abstract partial class AssistantBase : AssistantLowerBase wher return string.Empty; } + catch (PromptInjectionBlockedException e) + { + this.Logger.LogWarning(e, "Blocked prompt-injection content before sending assistant request for '{AssistantTitle}'.", this.Title); + + if (this.ResultingContentBlock is not null && string.IsNullOrWhiteSpace(aiText.Text)) + { + this.ChatThread?.Blocks.Remove(this.ResultingContentBlock); + this.ResultingContentBlock = null; + } + + return string.Empty; + } finally { this.IsProcessing = this.assistantSessionId is not null && (this.AssistantSessionService.TryGetSnapshot(this.assistantSessionKey)?.IsActive ?? false); diff --git a/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs b/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs index 0fed4451..5b674b1f 100644 --- a/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs +++ b/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs @@ -716,7 +716,7 @@ public partial class DocumentAnalysisAssistant : AssistantBaseCore