Runtime/AgentCore/Tools/ToolTracker.cs

using System;
using System.Collections.Generic;
namespace OnDeviceAgent.AgentCore
{

    public sealed class ToolTracker
    {
        readonly object m_Lock = new object();
        readonly HashSet<string> m_UsedToolNames = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
        readonly AgentToolGuardrailController m_Guardrail;
        readonly AgentTraceStore m_TraceStore;
        readonly bool m_GuardrailsEnabled;

        public ToolTracker(AgentToolGuardrailController guardrails, AgentTraceStore traceStore, bool guardrailsEnabled)
        {
            if (guardrails == null)
                throw new ArgumentNullException("guardrails");
            if (traceStore == null)
                throw new ArgumentNullException("traceStore");

            m_Guardrail = guardrails;
            m_TraceStore = traceStore;
            m_GuardrailsEnabled = guardrailsEnabled;
        }

        public IReadOnlyCollection<string> UsedToolNames
        {
            get
            {
                lock (m_Lock)
                    return new List<string>(m_UsedToolNames);
            }
        }

        public bool TryBegin(string toolName, string message, string argumentsSignature, out string blockedResult)
        {
            blockedResult = string.Empty;

            lock (m_Lock)
            {
                var traceName = toolName ?? string.Empty;
                if (!string.IsNullOrWhiteSpace(traceName))
                    m_UsedToolNames.Add(traceName);
            }

            m_TraceStore.RecordToolCall(toolName, message);

            if (!m_GuardrailsEnabled)
                return true;

            if (m_Guardrail.TryBeginToolCall(toolName, argumentsSignature, out blockedResult))
                return true;

            m_TraceStore.RecordEvent("guardrail", blockedResult);
            return false;
        }

        public string Finish(string toolName, string argumentsSignature, string result)
        {
            var finalResult = result ?? string.Empty;

            if (m_GuardrailsEnabled)
            {
                bool shouldWarn;
                finalResult = m_Guardrail.InspectToolResult(toolName, finalResult, out shouldWarn);
                if (shouldWarn)
                    m_TraceStore.RecordEvent("guardrail", "Repeated result from " + toolName + ".");
            }

            m_TraceStore.RecordToolResult(toolName, finalResult);
            return finalResult;
        }

        public void ResetForRun()
        {
            lock (m_Lock)
                m_UsedToolNames.Clear();

            m_Guardrail.ResetForRun();
        }
    }
}