tool_executor.hpp

#pragma once

#include "../provider/llm_provider.hpp"
#include "diff_utils.hpp"
#include "question_policy.hpp"

#include <nlohmann/json.hpp>

#include <string>
#include <vector>
#include <map>
#include <functional>
#include <atomic>
#include <optional>
#include <utility>
#include <mutex>
#include <unordered_set>

namespace acecode {

class SessionManager;
class SkillRegistry;
class ToolExecutor;

// Origin of a registered tool. MCP tools are grouped separately in the system
// prompt so the LLM can distinguish internal versus external capabilities.
enum class ToolSource {
    Builtin = 0,
    Mcp = 1,
};

// Session-local expert policy. Missing members inherit all currently
// registered/global capabilities, while engaged empty sets allow none.
// Built-in tool IDs and MCP server IDs remain separate namespaces.
struct ToolCapabilityPolicy {
    std::optional<std::unordered_set<std::string>> builtin_tools;
    std::optional<std::unordered_set<std::string>> mcp_servers;
};

// Structured summary used by tool-result UIs to render a single-line row
// (icon + verb + object + dot-separated metrics). Tool implementations may
// provide a domain-specific summary; the execution boundary supplies a generic
// summary when they do not.
struct ToolSummary {
    std::string verb;     // "Ran" / "Read" / "Wrote" / "Created" / "Edited" ...
    std::string object;   // file path or command preview
    std::vector<std::pair<std::string, std::string>> metrics; // ordered
    std::string icon;     // short glyph (may be Unicode or ASCII fallback)
};

// Result of a tool execution
struct ToolResult {
    std::string output;
    bool success = true;
    std::optional<ToolSummary> summary; // always populated at the execution boundary
    // Optional UI/persistence metadata. This is never part of provider-visible
    // text; AgentLoop stores it on the ChatMessage and web lifecycle payloads.
    nlohmann::json metadata = nlohmann::json::object();
    // Optional user-role prompt to append after this tool result. This is used
    // for progressive capability disclosure that should affect only the active
    // conversation after a tool is explicitly opened, rather than the global
    // cacheable system prompt.
    std::optional<std::string> post_user_prompt;
    std::string post_user_prompt_display_text;
    // 结构化 diff hunk。file_edit / file_write 在产生 unified diff 文本的同时
    // 填充这个字段;TUI 用它做彩色带行号 gutter 的渲染。
    // 运行时字段 —— 不写入 session JSONL(由 session_serializer 的 allowlist
    // 天然挡住;新加字段时如果不加进白名单就不会被序列化)。
    std::optional<std::vector<DiffHunk>> hunks;
    // Structured output attachments produced by a tool. Items are either stored
    // AttachmentRecord JSON objects or pre-materialization descriptors such as
    // {name,mime_type,data_url} / {name,mime_type,path}. AgentLoop materializes
    // descriptors into session attachments before events and JSONL persistence.
    nlohmann::json attachments = nlohmann::json::array();
    std::vector<std::string> attachment_warnings;

    // Runtime-only terminal control. A tool that permanently removes the
    // calling session uses this to stop the current model turn after its
    // canonical tool result and terminal lifecycle events have been persisted.
    // Neither field is serialized into the provider-visible result.
    bool terminate_session_after_turn = false;
    std::function<void()> post_turn_action;

    bool has_attachments() const {
        return attachments.is_array() && !attachments.empty();
    }
};

// Build and attach the shared fallback used by built-in, MCP, unknown, denied,
// failed, and legacy tool calls that do not provide a domain-specific summary.
// Existing summaries are never overwritten.
ToolSummary build_fallback_tool_summary(
    const std::string& tool_name,
    const std::string& arguments_json);
void ensure_tool_summary(
    const std::string& tool_name,
    const std::string& arguments_json,
    ToolResult& result);

struct ScratchPathResolution {
    bool success = true;
    bool used_alias = false;
    std::string path;
    std::string error;
};

// Runtime context passed into a tool invocation. Optional: if left
// default-constructed, tools behave as if no streaming/abort is available.
// Populated by AgentLoop before each tool call so the tool can push
// interim output to the TUI and react to Esc-driven aborts.
struct ToolContext {
    // Session workspace cwd. Tools that support a default working directory
    // should prefer this over the daemon process cwd when their own arguments
    // omit a cwd/path.
    std::string cwd;

    // Called zero or more times with non-empty cleaned chunks (ANSI stripped,
    // UTF-8 boundary safe, carriage-return overwrites resolved). Only bash_tool
    // uses this currently — other tools return their output atomically.
    std::function<void(const std::string& chunk)> stream;
    // Non-owning pointer to AgentLoop::abort_requested_. Tools with long polling
    // loops must check this every iteration and terminate their subprocess /
    // work when it becomes true.
    const std::atomic<bool>* abort_flag = nullptr;
    // Optional file-checkpoint hook used by write tools. Tools call this after
    // validation succeeds and immediately before mutating a file so /rewind can
    // restore the pre-write state.
    std::function<void(const std::string& path)> track_file_write_before;

    // Per-session scratch directory for temporary helper files. AgentLoop
    // injects `.acecode/tmp/session-<id>` under the workspace when a session id
    // is available. Shell tools expose this as ACECODE_TMPDIR.
    std::string scratch_dir;

    // Whether file_path is inside the workspace-managed temporary root. The
    // root is derived from scratch_dir (its parent), so callers never duplicate
    // the `.acecode/tmp` path contract.
    bool is_workspace_scratch_path(const std::string& file_path) const;

    // Recognizes the explicit ACECODE_TMPDIR spellings accepted by shell and
    // structured file tools. This intentionally does not expand arbitrary
    // environment variables.
    static bool references_scratch_path_alias(const std::string& value);

    // Resolves a leading ACECODE_TMPDIR path component to scratch_dir. Invalid
    // placement, unavailable context, and parent traversal fail closed before
    // a file tool reaches filesystem APIs.
    ScratchPathResolution resolve_scratch_path_alias(
        const std::string& file_path) const;

    // Optional async channel for AskUserQuestion. Daemon path injects an impl
    // backed by AskUserQuestionPrompter; TUI path keeps it null and registers
    // the TUI-flavored AskUserQuestion factory which talks to TuiState directly.
    //
    // Wire format (nlohmann::json) — kept loose so this header doesn't pull
    // in session/ headers:
    //   in  questions_payload: array of {id, text, options:[{label, value}], multiSelect}
    //   out: { cancelled: bool,
    //          answers: [ { question_id, selected: [str], custom_text: str } ] }
    // Empty function = AskUserQuestion tool returns the rejected ToolResult.
    std::function<nlohmann::json(const nlohmann::json& questions_payload)> ask_user_questions;

    // Per-session state injected by AgentLoop. Goal tools use this instead of
    // binding to one SessionManager at process-wide tool registration time.
    SessionManager* session_manager = nullptr;

    // Effective per-session Skill registry. Skill tools prefer this over
    // rebuilding a workspace-only registry, which preserves expert-bundled
    // Skill isolation and precedence.
    const SkillRegistry* skill_registry = nullptr;

    // AgentLoop sets this so bash can hand the full output to the
    // tool-result budget layer. Standalone tool callers keep the legacy
    // 100KB inline cap unless they explicitly opt in.
    bool preserve_full_output = false;

    std::function<void()> account_goal_usage;
    std::function<void(const nlohmann::json& goal_payload)> emit_goal_updated;
    std::function<void(const std::string& session_id)> emit_goal_cleared;
    std::function<void(const nlohmann::json& todo_payload)> emit_todo_updated;

    // Goal 运行探针(AgentLoop 注入)。true = 当前会话(或父会话)
    // 有 Active goal 且非 Plan mode。工具权限确认自动放行;
    // AskUserQuestion 正常弹 UI，但固定 30 秒后自动采纳推荐项。
    // 空函数 = 正常交互模式。
    std::function<bool()> goal_unattended_active;

    // AskUserQuestion 应答策略探针(AgentLoop 注入,模式同
    // goal_unattended_active)。返回 resolve_question_policy 的解析结果,
    // permission mode 不参与提问策略解析。空函数 = Ask(独立调用
    // ToolExecutor 的旧行为)。active goal 在工具入口覆盖为 Timeout(30)。
    std::function<ResolvedQuestionPolicy()> question_policy;

    // Plan-mode tools use these callbacks to mutate the active AgentLoop's
    // permission state. They are callbacks rather than direct PermissionManager
    // references so the tool layer stays independent of the TUI/daemon runtime.
    std::function<std::string()> current_permission_mode;
    std::function<std::string()> enter_plan_mode;
    std::function<std::string()> exit_plan_mode;

    // Worktree 工具回调(AgentLoop 注入):把会话工作目录切到 new_cwd,
    // 并以新根重建路径校验器。enter_worktree / exit_worktree 用它在
    // worktree 与原目录之间切换;空函数 = 当前 runtime 不支持切换
    // (独立 ToolExecutor 调用),工具会拒绝执行。
    std::function<void(const std::string& new_cwd)> switch_session_cwd;

    // Runtime access to the active executor. Tools that intentionally change
    // the available tool set can use this to register additional tools for the
    // next model request.
    ToolExecutor* tool_executor = nullptr;

    // Optional session-local expert scope. ToolExecutor checks this again at
    // the execution boundary so replayed/model-produced calls cannot bypass
    // provider schema filtering.
    std::optional<ToolCapabilityPolicy> capability_policy;

    // 当前 active 模型的身份与视觉能力(AgentLoop 注入)。vision_analyze 用它把
    // "当前模型"从候选视觉模型里剔除:主模型自己带 vision 标签时,子调用很容易
    // 又挑中同一个模型,变成绕一圈用同一个模型看同一张图(实测会话
    // 20260830-024351-9599 白烧了约 5k token)。空串 = 未接线,此时不做剔除以
    // 维持旧行为。
    std::string active_provider_name;
    std::string active_model_id;
    // 与 LlmProvider::supports_vision 同口径,默认 fail-open。
    bool active_model_can_read_images = true;
};

// UI-only metadata contract: a successful structured file change under the
// workspace scratch root must not contribute to the per-turn "modified files"
// summary. The tool row and its diff remain available.
inline constexpr const char* kExcludeFromTurnChangeSummaryMetadata =
    "exclude_from_turn_change_summary";

void mark_workspace_scratch_change(ToolResult& result, const ToolContext& ctx);

// A registered tool implementation. The execute function takes a ToolContext —
// tools that don't need streaming simply ignore it.
struct ToolImpl {
    ToolDef definition;
    std::function<ToolResult(const std::string& arguments_json, const ToolContext& ctx)> execute;
    bool is_read_only = false; // Read-only tools are auto-approved without user confirmation
    ToolSource source = ToolSource::Builtin;
    // Exact owning MCP server ID. Empty for built-ins. This metadata is never
    // inferred from a qualified tool name.
    std::string source_owner;
};

struct RegisteredToolInfo {
    ToolDef definition;
    bool is_read_only = false;
    ToolSource source = ToolSource::Builtin;
    std::string source_owner;
};

class ToolExecutor {
public:
    // Register or refresh a tool only when an existing entry has the same
    // source/owner identity. This prevents one MCP server from overwriting a
    // different server (or a built-in) on a qualified-name collision.
    bool register_tool(const ToolImpl& tool);

    // Remove a tool by name. When expected_source_owner is present, removal is
    // allowed only for the exact (name, owner) pair.
    bool unregister_tool(
        const std::string& name,
        std::optional<std::string> expected_source_owner = std::nullopt);

    // Get all tool definitions for inclusion in API requests
    std::vector<ToolDef> get_tool_definitions(
        const ToolCapabilityPolicy* policy = nullptr) const;

    // Get tool definitions filtered by source (built-in vs MCP).
    std::vector<ToolDef> get_tool_definitions_by_source(
        ToolSource source,
        const ToolCapabilityPolicy* policy = nullptr) const;

    // Get tool definitions translated to the public names exposed to models.
    // Internal callers should continue using the native definition methods.
    std::vector<ToolDef> get_model_tool_definitions(
        const ToolCapabilityPolicy* policy = nullptr) const;

    std::vector<ToolDef> get_model_tool_definitions_by_source(
        ToolSource source,
        const ToolCapabilityPolicy* policy = nullptr) const;

    // Accept an exact registered native name first, then resolve a compatible
    // public alias only when its native handler is registered.
    std::string resolve_model_tool_name_to_native(
        const std::string& model_name) const;

    // Sanitized registration metadata for runtime-backed capability catalogs.
    std::vector<RegisteredToolInfo> get_registered_tools() const;

    // Central policy predicate shared by schema filtering and AgentLoop's
    // pre-permission rejection path.
    bool is_allowed(const std::string& name,
                    const ToolCapabilityPolicy* policy) const;

    // Returns true only for a registered tool excluded by the policy. Unknown
    // tool names remain unknown instead of being mislabeled as policy denials.
    bool is_denied_by_policy(const std::string& name,
                             const ToolCapabilityPolicy* policy) const;

    // Execute a tool call and return the result. Legacy overload — no streaming,
    // no abort flag. Delegates to the ctx overload with a default context.
    ToolResult execute(const std::string& tool_name, const std::string& arguments_json) const;

    // Execute with a ToolContext. Tools that support streaming will call
    // ctx.stream() as chunks arrive; tools that support cancellation will
    // poll ctx.abort_flag.
    ToolResult execute(const std::string& tool_name, const std::string& arguments_json,
                       const ToolContext& ctx) const;

    // Check if a tool is registered
    bool has_tool(const std::string& name) const;

    // Check if a tool is read-only (auto-approved)
    bool is_read_only(const std::string& name) const;

    // Generate a formatted description of all registered tools for system prompt
    std::string generate_tools_prompt(
        const ToolCapabilityPolicy* policy = nullptr) const;

    // Format a tool result into a ChatMessage suitable for the messages array
    static ChatMessage format_tool_result(const std::string& tool_call_id, const ToolResult& result);

    // Format an assistant message that includes tool calls (from the API response)
    static ChatMessage format_assistant_tool_calls(const ChatResponse& response);

    // Build a compact one-line preview for a tool_call row. For bash takes the
    // command's first 60 chars; for file_read/file_write/file_edit takes the
    // file_path (tail-truncated to 40 chars); other tools return an empty
    // string so the TUI falls back to the legacy `[Tool: X] {JSON}` format.
    static std::string build_tool_call_preview(const std::string& tool_name,
                                               const std::string& arguments_json);

private:
    std::map<std::string, ToolImpl> tools_;
    mutable std::mutex tools_mu_;
};

} // namespace acecode