Let tool results restrict a chat to self-hosted providers

app/MindWork AI Studio/Chat/ChatThread.cs

@@ -106,6 +106,26 @@ public sealed record ChatThread
             this.RequiredProviderConfidence = minimumProviderConfidence;
     }
 
+    /// <summary>
+    /// Tightens the data security of this chat to what the data brought in demands, and never
+    /// loosens it.
+    /// </summary>
+    /// <remarks>
+    /// Data which may only be used with self-hosted providers keeps the chat restricted to them,
+    /// no matter what comes in later: the data was seen by this chat. Data which may be used with
+    /// any provider marks the chat as one which holds data of a data source, while a restriction set
+    /// earlier stays. NOT_SPECIFIED demands nothing and changes nothing.<br/><br/>
+    /// Shared by the RAG process and by tools which search the data sources, so both tighten a chat
+    /// the same way.
+    /// </remarks>
+    /// <param name="dataSecurity">What the data brought in demands.</param>
+    public void RequireDataSecurity(DataSourceSecurity dataSecurity) => this.DataSecurity = (this.DataSecurity, dataSecurity) switch
+    {
+        (DataSourceSecurity.SELF_HOSTED, _) or (_, DataSourceSecurity.SELF_HOSTED) => DataSourceSecurity.SELF_HOSTED,
+        (_, DataSourceSecurity.ALLOW_ANY) => DataSourceSecurity.ALLOW_ANY,
+        _ => this.DataSecurity,
+    };
+
     /// <summary>
     /// The name of the chat thread. Usually generated by an AI model or manually edited by the user.
     /// </summary>

app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs

@@ -136,48 +136,15 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
                 
                 //
                 // Update the data security of the chat thread. We consider the current data security
-                // of the chat thread and the data security of the selected data sources:
+                // of the chat thread and the data security of the selected data sources: at least
+                // one data source with a SELF_HOSTED policy restricts the chat to self-hosted
+                // providers. A restriction set earlier stays either way, because the thread might
+                // already contain data from a data source with a SELF_HOSTED policy:
                 //
                 var dataSecurityRestrictedToSelfHosted = selectedDataSources
                     .OfType<IExternalDataSource>()
                     .Any(dataSource => dataSource.SecurityPolicy is DataSourceSecurity.SELF_HOSTED);
-                chatThread.DataSecurity = dataSecurityRestrictedToSelfHosted switch
+                chatThread.RequireDataSecurity(dataSecurityRestrictedToSelfHosted ? DataSourceSecurity.SELF_HOSTED : DataSourceSecurity.ALLOW_ANY);
-                {
-                    //
-                    //
-                    // Case: the data sources which are selected have a security policy
-                    // of SELF_HOSTED (at least one data source).
-                    //
-                    // When the policy was already set to ALLOW_ANY, we restrict it
-                    // to SELF_HOSTED.
-                    //
-                    true => DataSourceSecurity.SELF_HOSTED,
-                    
-                    //
-                    // Case: the data sources which are selected have a security policy
-                    // of ALLOW_ANY (none of the data sources has a SELF_HOSTED policy).
-                    //
-                    // When the policy was already set to SELF_HOSTED, we must keep that.
-                    //
-                    false => chatThread.DataSecurity switch
-                    {
-                        //
-                        // When the policy was not specified yet, we set it to ALLOW_ANY.
-                        //
-                        DataSourceSecurity.NOT_SPECIFIED => DataSourceSecurity.ALLOW_ANY,
-                        DataSourceSecurity.ALLOW_ANY => DataSourceSecurity.ALLOW_ANY,
-                        
-                        //
-                        // When the policy was already set to SELF_HOSTED, we must keep that.
-                        // This is important since the thread might already contain data
-                        // from a data source with a SELF_HOSTED policy.
-                        //
-                        DataSourceSecurity.SELF_HOSTED => DataSourceSecurity.SELF_HOSTED,
-                        
-                        // Default case: we use the current data security of the chat thread.
-                        _ => chatThread.DataSecurity,
-                    }
-                };
                 
                 if (previousDataSecurity != chatThread.DataSecurity)
                     LOGGER.LogInformation($"The data security of the chat thread was updated from '{previousDataSecurity}' to '{chatThread.DataSecurity}'.");

app/MindWork AI Studio/Tools/ToolCallingSystem/Harness/ToolCallingLoop.cs

@@ -223,17 +223,19 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
                     }
 
                     toolCallCount++;
-                    var (toolContent, trace, requiredProviderConfidence, sources) = await context.ToolExecutor.ExecuteAsync(
+                    var (toolContent, trace, requiredProviderConfidence, requiredDataSecurity, sources) = await context.ToolExecutor.ExecuteAsync(
                         call.CallId,
                         call.ToolName,
                         call.ArgumentsJson,
                         context.RunnableTools,
                         context.Provider,
+                        context.ChatThread,
                         toolCallCount,
                         token);
 
                     toolResultCharacterCount += toolContent.Length;
                     context.ChatThread.RequireProviderConfidence(requiredProviderConfidence);
+                    context.ChatThread.RequireDataSecurity(requiredDataSecurity);
                     toolSources.MergeSources(sources);
                     await context.AddToolInvocationAsync(trace);
 

app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionContext.cs

@@ -1,3 +1,4 @@
+using AIStudio.Chat;
 using AIStudio.Provider;
 using AIStudio.Settings;
 
@@ -7,6 +8,16 @@ public sealed class ToolExecutionContext
 {
     public required ToolDefinition Definition { get; init; }
 
+    /// <summary>
+    /// The chat the call was made in.
+    /// </summary>
+    /// <remarks>
+    /// For a tool which works with what the chat was set up with, such as Semantic Search with the
+    /// data sources the user picked for it. A tool reads it; what the chat has to keep because of
+    /// the result goes back through the ToolExecutionResult instead.
+    /// </remarks>
+    public required ChatThread ChatThread { get; init; }
+
     public string ToolCallId { get; init; } = string.Empty;
 
     public required SettingsManager SettingsManager { get; init; }

app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionResult.cs

@@ -1,6 +1,7 @@
 using System.Text.Json.Nodes;
 
 using AIStudio.Provider;
+using AIStudio.Settings.DataModel;
 
 namespace AIStudio.Tools.ToolCallingSystem;
 
@@ -14,6 +15,17 @@ public sealed class ToolExecutionResult
 
     public ConfidenceLevel RequiredProviderConfidence { get; init; } = ConfidenceLevel.NONE;
 
+    /// <summary>
+    /// The data security the chat has to keep from now on, because of what this result brings in.
+    /// </summary>
+    /// <remarks>
+    /// The other axis next to RequiredProviderConfidence. A data source which may only be used with
+    /// self-hosted providers says so here, and the chat then refuses every other provider from now
+    /// on, see ChatThread.RequireDataSecurity. Left at NOT_SPECIFIED, the result says nothing about
+    /// it, and the chat stays as it was.
+    /// </remarks>
+    public DataSourceSecurity RequiredDataSecurity { get; init; } = DataSourceSecurity.NOT_SPECIFIED;
+
     public string ToModelContent()
     {
         if (this.JsonContent is not null)

app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutor.cs

@@ -1,8 +1,10 @@
 using System.Diagnostics;
 using System.Text.Json;
 
+using AIStudio.Chat;
 using AIStudio.Provider;
 using AIStudio.Settings;
+using AIStudio.Settings.DataModel;
 
 namespace AIStudio.Tools.ToolCallingSystem;
 
@@ -46,12 +48,13 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge
         }
     }
 
-    public async Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, IReadOnlyList<Source> Sources)> ExecuteAsync(
+    public async Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, DataSourceSecurity RequiredDataSecurity, IReadOnlyList<Source> Sources)> ExecuteAsync(
         string toolCallId,
         string toolName,
         string argumentsJson,
         IReadOnlyList<(ToolDefinition Definition, IToolImplementation Implementation)> runnableTools,
         IProvider provider,
+        ChatThread chatThread,
         int order,
         CancellationToken token = default)
     {
@@ -92,7 +95,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge
                 StatusMessage = "Tool is not available in the current context.",
                 Arguments = formattedArguments,
                 Result = error,
-            }, ConfidenceLevel.NONE, []);
+            }, ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED, []);
         }
 
         var definition = runnableTool.Definition;
@@ -105,6 +108,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge
             var result = await implementation.ExecuteAsync(document.RootElement, new ToolExecutionContext
             {
                 Definition = definition,
+                ChatThread = chatThread,
                 ToolCallId = toolCallId,
                 SettingsManager = settingsManager,
                 SettingsValues = settingsValues,
@@ -128,7 +132,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge
                 JsonResult = result.JsonContent,
             };
 
-            return (resultModelContent, toolInvocationTrace, result.RequiredProviderConfidence, result.Sources);
+            return (resultModelContent, toolInvocationTrace, result.RequiredProviderConfidence, result.RequiredDataSecurity, result.Sources);
         }
         catch (OperationCanceledException) when (token.IsCancellationRequested)
         {
@@ -152,7 +156,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge
                 Result = exception.Message,
             };
 
-            return (exception.Message, toolInvocationTrace, ConfidenceLevel.NONE, []);
+            return (exception.Message, toolInvocationTrace, ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED, []);
         }
         catch (Exception exception)
         {
@@ -172,7 +176,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge
                 Result = error,
             };
 
-            return (error, toolInvocationTrace, ConfidenceLevel.NONE, []);
+            return (error, toolInvocationTrace, ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED, []);
         }
     }
 

app/Tests/Chat/ChatThreadDataSecurityTests.cs

@@ -0,0 +1,52 @@
+using AIStudio.Chat;
+using AIStudio.Settings.DataModel;
+
+namespace AIStudio.Tests.Chat;
+
+/// <summary>
+/// Checks how the data a chat has seen tightens the providers which may continue it.
+/// </summary>
+/// <remarks>
+/// A chat which once held data for self-hosted providers only must never be sent to any other
+/// provider again, whatever it brings in afterwards. The RAG process and Semantic Search both go
+/// through the same rule, so every combination of what a chat holds and what arrives is checked.
+/// </remarks>
+[TestFixture]
+public sealed class ChatThreadDataSecurityTests
+{
+    [TestCase(DataSourceSecurity.NOT_SPECIFIED, DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED)]
+    [TestCase(DataSourceSecurity.ALLOW_ANY, DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED)]
+    [TestCase(DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED)]
+    public void DataForSelfHostedProvidersOnlyRestrictsTheChat(DataSourceSecurity held, DataSourceSecurity arriving, DataSourceSecurity expected)
+    {
+        Assert.That(Tightened(held, arriving), Is.EqualTo(expected));
+    }
+
+    [TestCase(DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.ALLOW_ANY)]
+    [TestCase(DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.NOT_SPECIFIED)]
+    public void ARestrictionStays(DataSourceSecurity held, DataSourceSecurity arriving)
+    {
+        Assert.That(Tightened(held, arriving), Is.EqualTo(DataSourceSecurity.SELF_HOSTED), "The restricted data was seen by this chat. What arrives later cannot undo that.");
+    }
+
+    [TestCase(DataSourceSecurity.NOT_SPECIFIED)]
+    [TestCase(DataSourceSecurity.ALLOW_ANY)]
+    public void DataForAnyProviderMarksTheChat(DataSourceSecurity held)
+    {
+        Assert.That(Tightened(held, DataSourceSecurity.ALLOW_ANY), Is.EqualTo(DataSourceSecurity.ALLOW_ANY));
+    }
+
+    [TestCase(DataSourceSecurity.NOT_SPECIFIED)]
+    [TestCase(DataSourceSecurity.ALLOW_ANY)]
+    public void ResultsWhichDemandNothingChangeNothing(DataSourceSecurity held)
+    {
+        Assert.That(Tightened(held, DataSourceSecurity.NOT_SPECIFIED), Is.EqualTo(held), "A web search, say, says nothing about data sources and must leave the chat as it was.");
+    }
+
+    private static DataSourceSecurity Tightened(DataSourceSecurity held, DataSourceSecurity arriving)
+    {
+        var thread = new ChatThread { DataSecurity = held };
+        thread.RequireDataSecurity(arriving);
+        return thread.DataSecurity;
+    }
+}

app/Tests/Tools/ToolCalling/ToolExecutorTests.cs

@@ -0,0 +1,90 @@
+using AIStudio.Chat;
+using AIStudio.Provider;
+using AIStudio.Settings.DataModel;
+using AIStudio.Tools.ToolCallingSystem;
+
+using Microsoft.Extensions.Logging.Abstractions;
+
+namespace AIStudio.Tests.Tools.ToolCalling;
+
+/// <summary>
+/// Checks what the tool executor hands to a tool and what it hands back to the loop.
+/// </summary>
+/// <remarks>
+/// What a result demands of the chat has to reach the loop, which tightens the chat with it: a
+/// result from a data source for self-hosted providers only that got lost on the way would let the
+/// next message go to a cloud provider. A call which brought nothing in must demand nothing.
+/// </remarks>
+[TestFixture]
+[NonParallelizable]
+public sealed class ToolExecutorTests : ToolRegistryTestBase
+{
+    [Test]
+    public async Task WhatAResultDemandsReachesTheLoop()
+    {
+        var tool = new TestTool(Definition(), execute: _ => new ToolExecutionResult
+        {
+            TextContent = "A passage from the handbook.",
+            RequiredProviderConfidence = ConfidenceLevel.HIGH,
+            RequiredDataSecurity = DataSourceSecurity.SELF_HOSTED,
+        });
+
+        var (_, _, requiredProviderConfidence, requiredDataSecurity, _) = await this.Execute(tool, new ChatThread());
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(requiredDataSecurity, Is.EqualTo(DataSourceSecurity.SELF_HOSTED));
+            Assert.That(requiredProviderConfidence, Is.EqualTo(ConfidenceLevel.HIGH));
+        });
+    }
+
+    [Test]
+    public async Task TheToolSeesTheChatOfTheCall()
+    {
+        ChatThread? seenThread = null;
+        var tool = new TestTool(Definition(), execute: context =>
+        {
+            seenThread = context.ChatThread;
+            return new ToolExecutionResult();
+        });
+        var thread = new ChatThread();
+
+        await this.Execute(tool, thread);
+
+        Assert.That(seenThread, Is.SameAs(thread), "Semantic Search searches the data sources picked for this very chat.");
+    }
+
+    [Test]
+    public async Task ABlockedCallDemandsNothing()
+    {
+        var tool = new TestTool(Definition(), execute: _ => throw new ToolExecutionBlockedException("The data source is not available to this provider."));
+
+        var (_, trace, _, requiredDataSecurity, _) = await this.Execute(tool, new ChatThread());
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(trace.Status, Is.EqualTo(ToolInvocationTraceStatus.BLOCKED));
+            Assert.That(requiredDataSecurity, Is.EqualTo(DataSourceSecurity.NOT_SPECIFIED), "Nothing reached the model, so there is nothing the chat has to keep.");
+        });
+    }
+
+    [Test]
+    public async Task AFailedCallDemandsNothing()
+    {
+        var tool = new TestTool(Definition(), execute: _ => throw new InvalidOperationException("The index could not be read."));
+
+        var (_, trace, _, requiredDataSecurity, _) = await this.Execute(tool, new ChatThread());
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(trace.Status, Is.EqualTo(ToolInvocationTraceStatus.ERROR));
+            Assert.That(requiredDataSecurity, Is.EqualTo(DataSourceSecurity.NOT_SPECIFIED), "Nothing reached the model, so there is nothing the chat has to keep.");
+        });
+    }
+
+    private Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, DataSourceSecurity RequiredDataSecurity, IReadOnlyList<AIStudio.Tools.Source> Sources)> Execute(TestTool tool, ChatThread thread)
+    {
+        var executor = new ToolExecutor(this.CreateToolSettingsService(), NullLogger<ToolExecutor>.Instance);
+        return executor.ExecuteAsync("call-1", TOOL_ID, "{}", [(tool.GetDefinition(), tool)], new NoProvider(), thread, order: 1);
+    }
+}

app/Tests/Tools/ToolCalling/ToolRegistryTestBase.cs

@@ -60,11 +60,9 @@ public abstract class ToolRegistryTestBase
         this.rustService.Dispose();
     }
 
-    protected ToolRegistry CreateRegistry(params TestTool[] tools)
+    protected ToolRegistry CreateRegistry(params TestTool[] tools) => new(tools, [new CodeToolDefinitionSource(tools)], this.SettingsManager, this.CreateToolSettingsService(), NullLogger<ToolRegistry>.Instance);
-    {
+
-        var toolSettingsService = new ToolSettingsService(this.SettingsManager, this.rustService, NullLogger<ToolSettingsService>.Instance);
+    protected ToolSettingsService CreateToolSettingsService() => new(this.SettingsManager, this.rustService, NullLogger<ToolSettingsService>.Instance);
-        return new ToolRegistry(tools, [new CodeToolDefinitionSource(tools)], this.SettingsManager, toolSettingsService, NullLogger<ToolRegistry>.Instance);
-    }
 
     protected ToolResolutionContext ContextFor(AIStudio.Settings.Provider provider) => new()
     {
@@ -104,11 +102,12 @@ public abstract class ToolRegistryTestBase
     };
 
     /// <summary>
-    /// A tool which does nothing, and offers what it is told to.
+    /// A tool which offers and returns what it is told to.
     /// </summary>
     /// <param name="definition">What the tool is.</param>
     /// <param name="resolve">What it offers per request; when left out, its function as defined.</param>
-    protected sealed class TestTool(ToolDefinition definition, Func<ToolDefinition, ToolFunctionDefinition?>? resolve = null) : IToolImplementation
+    /// <param name="execute">What a call returns; when left out, an empty result.</param>
+    protected sealed class TestTool(ToolDefinition definition, Func<ToolDefinition, ToolFunctionDefinition?>? resolve = null, Func<ToolExecutionContext, ToolExecutionResult>? execute = null) : IToolImplementation
     {
         public int ResolveCount { get; private set; }
 
@@ -124,6 +123,6 @@ public abstract class ToolRegistryTestBase
 
         public IReadOnlySet<string> SensitiveTraceArgumentNames { get; } = new HashSet<string>(StringComparer.Ordinal);
 
-        public Task<ToolExecutionResult> ExecuteAsync(JsonElement arguments, ToolExecutionContext context, CancellationToken token = default) => Task.FromResult(new ToolExecutionResult());
+        public Task<ToolExecutionResult> ExecuteAsync(JsonElement arguments, ToolExecutionContext context, CancellationToken token = default) => Task.FromResult(execute is null ? new ToolExecutionResult() : execute(context));
     }
 }