Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,7 @@ add_library(flapi-lib STATIC
src/rate_limit_key_builder.cpp
src/rate_limit_middleware.cpp
src/mcp_authorization_policy.cpp
src/mcp_dry_run.cpp
src/prepared_template_rewriter.cpp
src/route_translator.cpp
src/security_auditor.cpp
Expand Down
34 changes: 34 additions & 0 deletions src/include/mcp_dry_run.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
#pragma once

#include <crow/json.h>
#include <map>
#include <string>

namespace flapi {

// Helpers for W2.2 dry-run / shadow mode. The model is:
// 1. Caller sends `_dryRun: true` alongside the normal tool arguments.
// 2. MCPToolHandler peels the flag off (no validation impact downstream).
// 3. After auth + argument validation + template rendering, the handler
// returns the rendered SQL instead of executing it.
//
// MCPDryRun groups the flag-stripping helper and the result formatter so
// they can be unit-tested in isolation without spinning up a server.
class MCPDryRun {
public:
static constexpr const char* kFlagKey = "_dryRun";

// If `arguments` contains the reserved `_dryRun` key, strip it and return
// its boolean value. Non-boolean values are treated as false but still
// stripped, so a hostile caller can't smuggle the key into validation.
static bool extractFlag(crow::json::wvalue& arguments);

// Render the dry-run JSON payload returned to the agent in place of real
// query results. Always emits a `parameters` object (possibly empty) so
// callers don't need to special-case missing args.
static std::string formatResult(const std::string& tool_name,
const std::string& rendered_sql,
const std::map<std::string, std::string>& parameters);
};

} // namespace flapi
65 changes: 65 additions & 0 deletions src/mcp_dry_run.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
#include "mcp_dry_run.hpp"

namespace flapi {

bool MCPDryRun::extractFlag(crow::json::wvalue& arguments) {
if (arguments.t() != crow::json::type::Object) {
return false;
}
auto keys = arguments.keys();
bool present = false;
for (const auto& k : keys) {
if (k == kFlagKey) {
present = true;
break;
}
}
if (!present) {
return false;
}

// We have to round-trip via the rvalue type to read the value back, since
// crow::json::wvalue does not expose getters for individual children.
auto dumped = arguments.dump();
auto parsed = crow::json::load(dumped);
bool flag_value = false;
if (parsed && parsed.has(kFlagKey)) {
const auto& node = parsed[kFlagKey];
if (node.t() == crow::json::type::True) {
flag_value = true;
}
}

// Rebuild the wvalue without the reserved key so downstream validators
// never observe `_dryRun` as an unknown parameter.
crow::json::wvalue rebuilt;
if (parsed) {
for (const auto& key : parsed.keys()) {
if (key == kFlagKey) {
continue;
}
rebuilt[key] = parsed[key];
}
}
arguments = std::move(rebuilt);
return flag_value;
}

std::string MCPDryRun::formatResult(const std::string& tool_name,
const std::string& rendered_sql,
const std::map<std::string, std::string>& parameters) {
crow::json::wvalue payload;
payload["dry_run"] = true;
payload["tool_name"] = tool_name;
payload["rendered_sql"] = rendered_sql;

crow::json::wvalue params_obj = crow::json::wvalue::object();
for (const auto& [k, v] : parameters) {
params_obj[k] = v;
}
payload["parameters"] = std::move(params_obj);

return payload.dump();
}

} // namespace flapi
36 changes: 33 additions & 3 deletions src/mcp_tool_handler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
#include <sstream>
#include <algorithm>

#include "mcp_dry_run.hpp"

namespace flapi {

MCPToolHandler::MCPToolHandler(std::shared_ptr<DatabaseManager> db_manager,
Expand Down Expand Up @@ -74,14 +76,42 @@ MCPToolExecutionResult MCPToolHandler::executeTool(const MCPToolCallRequest& req
}
}

// Validate arguments
if (!validateToolArguments(request.tool_name, request.arguments)) {
// W2.2 dry-run: peel `_dryRun` off the arguments before validation so
// the reserved key never reaches the unknown-parameter check. A copy
// of the arguments is made because MCPToolCallRequest is const here.
crow::json::wvalue effective_arguments;
{
auto reparsed = crow::json::load(request.arguments.dump());
if (reparsed) {
effective_arguments = crow::json::wvalue(reparsed);
}
}
const bool is_dry_run = MCPDryRun::extractFlag(effective_arguments);

// Validate arguments (post-strip).
if (!validateToolArguments(request.tool_name, effective_arguments)) {
emit_audit("error:invalid_arguments", -1);
return createErrorResult("Invalid arguments for tool: " + request.tool_name);
}

// Prepare parameters for SQL template
std::map<std::string, std::string> params = prepareParameters(*endpoint_config, request.arguments);
std::map<std::string, std::string> params = prepareParameters(*endpoint_config, effective_arguments);

// W2.2 dry-run short-circuit: render the SQL via the existing template
// processor and return it without touching the database. Write tools
// honour dry-run the same way — no side effects, just the SQL that
// would have run.
if (is_dry_run) {
std::string rendered_sql = sql_processor->loadAndProcessTemplate(*endpoint_config, params);
std::string payload = MCPDryRun::formatResult(request.tool_name, rendered_sql, params);

std::unordered_map<std::string, std::string> metadata;
metadata["tool_name"] = request.tool_name;
metadata["dry_run"] = "true";
metadata["execution_time_ms"] = "0";

return createSuccessResult(payload, metadata);
}

// Check if this is a write operation
if (endpoint_config->operation.type == OperationConfig::Write) {
Expand Down
1 change: 1 addition & 0 deletions test/cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ add_executable(flapi_tests
https_config_test.cpp
cache_manager_test.cpp
mcp_authorization_policy_test.cpp
mcp_dry_run_test.cpp
mcp_prompt_handler_test.cpp
mcp_request_validator_test.cpp
password_hasher_test.cpp
Expand Down
97 changes: 97 additions & 0 deletions test/cpp/mcp_dry_run_test.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
#include <catch2/catch_test_macros.hpp>
#include <crow/json.h>

#include "mcp_dry_run.hpp"

namespace flapi {
namespace test {

TEST_CASE("MCPDryRun::extractFlag: missing key yields false and no change",
"[security][mcp][dryrun]") {
crow::json::wvalue args;
args["id"] = 42;

bool extracted = MCPDryRun::extractFlag(args);

REQUIRE_FALSE(extracted);
// The original argument must still be present.
auto dumped = args.dump();
REQUIRE(dumped.find("\"id\":42") != std::string::npos);
}

TEST_CASE("MCPDryRun::extractFlag: _dryRun=true is consumed and returns true",
"[security][mcp][dryrun]") {
crow::json::wvalue args;
args["id"] = 42;
args["_dryRun"] = true;

bool extracted = MCPDryRun::extractFlag(args);

REQUIRE(extracted);
auto dumped = args.dump();
// The flag must be stripped so downstream validators do not see it.
REQUIRE(dumped.find("_dryRun") == std::string::npos);
// Other arguments must survive untouched.
REQUIRE(dumped.find("\"id\":42") != std::string::npos);
}

TEST_CASE("MCPDryRun::extractFlag: _dryRun=false is consumed and returns false",
"[security][mcp][dryrun]") {
crow::json::wvalue args;
args["_dryRun"] = false;

bool extracted = MCPDryRun::extractFlag(args);

REQUIRE_FALSE(extracted);
auto dumped = args.dump();
REQUIRE(dumped.find("_dryRun") == std::string::npos);
}

TEST_CASE("MCPDryRun::extractFlag: non-boolean _dryRun is rejected and stripped",
"[security][mcp][dryrun]") {
// A string or numeric _dryRun is treated as not-set; we still strip the
// key so it never reaches the validator. This is conservative: only an
// explicit boolean true engages dry-run.
crow::json::wvalue args;
args["_dryRun"] = "yes";

bool extracted = MCPDryRun::extractFlag(args);

REQUIRE_FALSE(extracted);
auto dumped = args.dump();
REQUIRE(dumped.find("_dryRun") == std::string::npos);
}

TEST_CASE("MCPDryRun::formatResult: produces JSON with dry_run, tool, sql, params",
"[security][mcp][dryrun]") {
std::map<std::string, std::string> params = {
{"id", "42"},
{"region", "EU"},
};
std::string rendered = "SELECT * FROM customers WHERE id = 42 AND region = 'EU'";

std::string payload = MCPDryRun::formatResult("customer_lookup", rendered, params);

auto parsed = crow::json::load(payload);
REQUIRE(parsed);
REQUIRE(parsed["dry_run"].b() == true);
REQUIRE(parsed["tool_name"].s() == std::string("customer_lookup"));
REQUIRE(parsed["rendered_sql"].s() == rendered);
// Parameters must round-trip as a JSON object keyed by name.
REQUIRE(parsed["parameters"]["id"].s() == std::string("42"));
REQUIRE(parsed["parameters"]["region"].s() == std::string("EU"));
}

TEST_CASE("MCPDryRun::formatResult: empty parameter map still emits a parameters object",
"[security][mcp][dryrun]") {
std::string payload = MCPDryRun::formatResult(
"no_arg_tool", "SELECT 1", /*parameters=*/{});

auto parsed = crow::json::load(payload);
REQUIRE(parsed);
REQUIRE(parsed["dry_run"].b() == true);
REQUIRE(parsed.has("parameters"));
}

} // namespace test
} // namespace flapi
Loading