From 00b7252fbf54b38a950a7cb99cf52905bf9aacdb Mon Sep 17 00:00:00 2001 From: Brandon Wolf <126610820+WolfBrandon@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:49:43 -0400 Subject: [PATCH] feat(mcp): Add MCP budget update tool (#2908) * feat(mcp): Add MCP budget update tool Adds an update_budget assistant/MCP function so AI assistants can write monthly budgets: total budgeted spending, expected income, and per-category allocations in one transactional call. - Month resolution and slug format mirror get_budget (YYYY-MM or MMM-YYYY, custom month start respected); targeting a valid month with no budget row bootstraps it via Budget.find_or_bootstrap, same as the budgets UI. - Category allocations accept an exact (case-insensitive) name or id and go through BudgetCategory#update_budgeted_spending!, so subcategory writes keep the parent total in sync. - All writes in one call share a transaction: an invalid category rolls back a totals change from the same call. - Family-scoped like the budgets UI; amounts validated non-negative. Co-Authored-By: Claude Fable 5 * refactor(mcp): harden update_budget per review feedback - Extract shared month resolution into Assistant::Function::MonthResolvable so get_budget and update_budget can't drift on custom month starts - Run budget bootstrap inside the update transaction so a failed entry no longer leaves a newly created budget behind - Apply explicit parent amounts after subcategory syncs so results don't depend on the caller's array order - Reject non-finite amounts (NaN/Infinity) - Explain the synthetic Uncategorized bucket instead of a generic category-not-found error Co-Authored-By: Claude Fable 5 --------- Co-authored-by: Claude Fable 5 --- app/models/assistant.rb | 3 +- app/models/assistant/function/get_budget.rb | 25 +-- .../assistant/function/month_resolvable.rb | 29 +++ .../assistant/function/update_budget.rb | 181 +++++++++++++++ test/controllers/mcp_controller_test.rb | 25 +++ .../assistant/function/update_budget_test.rb | 211 ++++++++++++++++++ 6 files changed, 449 insertions(+), 25 deletions(-) create mode 100644 app/models/assistant/function/month_resolvable.rb create mode 100644 app/models/assistant/function/update_budget.rb create mode 100644 test/models/assistant/function/update_budget_test.rb diff --git a/app/models/assistant.rb b/app/models/assistant.rb index 380b2f1b6..1cb432d61 100644 --- a/app/models/assistant.rb +++ b/app/models/assistant.rb @@ -52,7 +52,8 @@ module Assistant Function::GetCategories, Function::CreateCategory, Function::UpdateCategory, - Function::UpdateTransaction + Function::UpdateTransaction, + Function::UpdateBudget ] classes += PREVIEW_FUNCTION_CLASSES if user&.preview_features_enabled? diff --git a/app/models/assistant/function/get_budget.rb b/app/models/assistant/function/get_budget.rb index 3cf95115b..3d2a3abf0 100644 --- a/app/models/assistant/function/get_budget.rb +++ b/app/models/assistant/function/get_budget.rb @@ -1,5 +1,6 @@ class Assistant::Function::GetBudget < Assistant::Function include ActiveSupport::NumberHelper + include Assistant::Function::MonthResolvable MAX_PRIOR_MONTHS = 11 @@ -157,30 +158,6 @@ class Assistant::Function::GetBudget < Assistant::Function "no_activity" end - def resolve_month_start(raw) - base = parse_month(raw) - return (base || Date.current).beginning_of_month unless family.uses_custom_month_start? - - # Match Budget.param_to_date for explicit slugs so the input round-trips with the response. - base ? Date.new(base.year, base.month, family.month_start_day) : family.custom_month_start_for(Date.current) - end - - def parse_month(raw) - return nil if raw.blank? - - # Date.strptime ignores trailing characters, so guard with strict anchors first. - fmt = case raw - when /\A\d{4}-\d{2}\z/ then "%Y-%m" - when /\A[A-Za-z]{3}-\d{4}\z/ then "%b-%Y" - end - - raise Assistant::Error, "Invalid month: #{raw}. Use YYYY-MM or MMM-YYYY." if fmt.nil? - - Date.strptime(raw, fmt) - rescue ArgumentError - raise Assistant::Error, "Invalid month: #{raw}. Use YYYY-MM or MMM-YYYY." - end - def shift_months(date, n) shifted = date >> n if family.uses_custom_month_start? diff --git a/app/models/assistant/function/month_resolvable.rb b/app/models/assistant/function/month_resolvable.rb new file mode 100644 index 000000000..420dc3fbe --- /dev/null +++ b/app/models/assistant/function/month_resolvable.rb @@ -0,0 +1,29 @@ +# Shared month handling for budget-facing assistant tools so month slugs +# round-trip between get_budget and update_budget, including families with a +# custom month start day. +module Assistant::Function::MonthResolvable + private + def resolve_month_start(raw) + base = parse_month(raw) + return (base || Date.current).beginning_of_month unless family.uses_custom_month_start? + + # Match Budget.param_to_date for explicit slugs so the input round-trips with the response. + base ? Date.new(base.year, base.month, family.month_start_day) : family.custom_month_start_for(Date.current) + end + + def parse_month(raw) + return nil if raw.blank? + + # Date.strptime ignores trailing characters, so guard with strict anchors first. + fmt = case raw + when /\A\d{4}-\d{2}\z/ then "%Y-%m" + when /\A[A-Za-z]{3}-\d{4}\z/ then "%b-%Y" + end + + raise Assistant::Error, "Invalid month: #{raw}. Use YYYY-MM or MMM-YYYY." if fmt.nil? + + Date.strptime(raw, fmt) + rescue ArgumentError + raise Assistant::Error, "Invalid month: #{raw}. Use YYYY-MM or MMM-YYYY." + end +end diff --git a/app/models/assistant/function/update_budget.rb b/app/models/assistant/function/update_budget.rb new file mode 100644 index 000000000..77d07ee52 --- /dev/null +++ b/app/models/assistant/function/update_budget.rb @@ -0,0 +1,181 @@ +class Assistant::Function::UpdateBudget < Assistant::Function + include Assistant::Function::MonthResolvable + + class << self + def name + "update_budget" + end + + def description + <<~INSTRUCTIONS + Updates the user's monthly budget: total budgeted spending, expected income, + and/or per-category budgeted amounts. + + Call get_budget first to see current amounts and exact category names. + Amounts are plain non-negative numbers in the family's currency. Only the + fields and categories you pass are changed. Setting a subcategory's amount + keeps its parent's total in sync automatically. The "Uncategorized" bucket + cannot be set directly — it is the unallocated remainder of total budgeted + spending. + + Parameters: + - `month` (optional): "YYYY-MM" or "MMM-YYYY". Defaults to the current month. + - `budgeted_spending` (optional): total planned spending for the month. + - `expected_income` (optional): expected income for the month. + - `categories` (optional): array of { category: , amount: }. + + At least one of budgeted_spending, expected_income, or categories is required. + + Example (set the total and two category allocations for August 2026): + + ``` + update_budget({ + month: "2026-08", + budgeted_spending: 6500, + categories: [ + { category: "Groceries", amount: 900 }, + { category: "Dining Out", amount: 250 } + ] + }) + ``` + INSTRUCTIONS + end + end + + def strict_mode? + false + end + + def params_schema + build_schema( + properties: { + month: { + type: "string", + description: "Target month in YYYY-MM or MMM-YYYY format. Defaults to the current month." + }, + budgeted_spending: { + type: "number", + minimum: 0, + description: "Total planned spending for the month, in the family currency." + }, + expected_income: { + type: "number", + minimum: 0, + description: "Expected income for the month, in the family currency." + }, + categories: { + type: "array", + description: "Per-category budget allocations to set.", + items: { + type: "object", + properties: { + category: { + type: "string", + description: "Category name (exact match, case-insensitive) or category id (use get_categories)." + }, + amount: { + type: "number", + minimum: 0, + description: "New budgeted amount for this category." + } + }, + required: [ "category", "amount" ], + additionalProperties: false + } + } + } + ) + end + + def call(params = {}) + category_changes = Array(params["categories"]) + unless params.key?("budgeted_spending") || params.key?("expected_income") || category_changes.any? + return error("no_changes", "Provide at least one of budgeted_spending, expected_income, or categories.") + end + + start_date = resolve_month_start(params["month"]) + unless Budget.budget_date_valid?(start_date, family: family) + return error("invalid_month", "No budget exists (or can be created) for that month — it is outside the valid budget range.") + end + + attrs = {} + attrs[:budgeted_spending] = parse_amount!(params["budgeted_spending"], "budgeted_spending") if params.key?("budgeted_spending") + attrs[:expected_income] = parse_amount!(params["expected_income"], "expected_income") if params.key?("expected_income") + + budget = nil + updated = [] + # Bootstrap and all writes share one transaction so a bad entry can't + # leave a newly created (or half-updated) budget behind. + Budget.transaction do + budget = Budget.find_or_bootstrap(family, start_date: start_date, user: user) + + budget.update!(attrs) if attrs.any? + + changes = category_changes.map do |change| + budget_category = find_budget_category!(budget, change.is_a?(Hash) ? change["category"] : nil) + [ budget_category, parse_amount!(change["amount"], "amount for '#{budget_category.name}'") ] + end + + # Subcategory updates sync their parent's total, so explicit parent + # amounts apply last to keep results independent of the array order. + subcategories, parents = changes.partition { |budget_category, _amount| budget_category.subcategory? } + (subcategories + parents).each do |budget_category, amount| + budget_category.update_budgeted_spending!(amount) + updated << { category: budget_category.name, budgeted_spending: format_money(budget_category.reload.budgeted_spending) } + end + end + + budget.reload + { + success: true, + month: budget.to_param, + totals: { + budgeted_spending: format_money(budget.budgeted_spending), + expected_income: format_money(budget.expected_income), + allocated_spending: format_money(budget.allocated_spending), + available_to_allocate: format_money(budget.available_to_allocate) + }, + updated_categories: updated, + message: "Budget for #{budget.start_date.strftime('%B %Y')} updated." + } + rescue Assistant::Error => e + error("invalid_params", e.message) + rescue ActiveRecord::RecordInvalid => e + error("validation_failed", e.record.errors.full_messages.join("; ")) + end + + private + def parse_amount!(raw, label) + value = Float(raw) + raise Assistant::Error, "#{label} must be a non-negative number." if !value.finite? || value.negative? + value + rescue ArgumentError, TypeError + raise Assistant::Error, "#{label} must be a non-negative number." + end + + def find_budget_category!(budget, ref) + ref = ref.to_s.strip + raise Assistant::Error, "Each categories entry needs a category name or id." if ref.blank? + + category = valid_uuid?(ref) ? family.categories.find_by(id: ref) : nil + category ||= family.categories.where("LOWER(name) = ?", ref.downcase).first + + if category.nil? + if Category.all_uncategorized_names.any? { |name| name.casecmp?(ref) } + raise Assistant::Error, "'#{ref}' is the unallocated remainder of budgeted_spending and cannot be set directly. Adjust budgeted_spending or category amounts instead." + end + raise Assistant::Error, "Category '#{ref}' not found. Use get_categories to list categories." + end + + budget.budget_categories.find_by(category_id: category.id) || + raise(Assistant::Error, "No budget row exists for category '#{category.name}' in #{budget.to_param}.") + end + + def format_money(value) + Money.new(value || 0, family.currency).format + end + + def error(key, message) + { success: false, error: key, message: message } + end +end diff --git a/test/controllers/mcp_controller_test.rb b/test/controllers/mcp_controller_test.rb index fdbcb7f22..aec44a8c8 100644 --- a/test/controllers/mcp_controller_test.rb +++ b/test/controllers/mcp_controller_test.rb @@ -289,6 +289,7 @@ class McpControllerTest < ActionDispatch::IntegrationTest assert_includes tool_names, "get_balance_sheet" assert_includes tool_names, "get_income_statement" assert_includes tool_names, "update_transaction" + assert_includes tool_names, "update_budget" # Each tool has required fields tools.each do |tool| @@ -463,6 +464,30 @@ class McpControllerTest < ActionDispatch::IntegrationTest end end + test "tools/call executes update_budget" do + with_mcp_env do + budget = budgets(:one) + + post "/mcp", params: jsonrpc_request("tools/call", { + name: "update_budget", + arguments: { + budgeted_spending: 6200, + expected_income: 8800 + } + }).to_json, headers: mcp_headers(@token) + + assert_response :ok + body = JSON.parse(response.body) + result = body["result"] + inner = JSON.parse(result["content"][0]["text"]) + + assert_equal true, inner["success"] + budget.reload + assert_equal 6200, budget.budgeted_spending + assert_equal 8800, budget.expected_income + end + end + test "tools/call wraps function errors as isError response" do with_mcp_env do # Force a function error by stubbing diff --git a/test/models/assistant/function/update_budget_test.rb b/test/models/assistant/function/update_budget_test.rb new file mode 100644 index 000000000..8eb985084 --- /dev/null +++ b/test/models/assistant/function/update_budget_test.rb @@ -0,0 +1,211 @@ +require "test_helper" + +class Assistant::Function::UpdateBudgetTest < ActiveSupport::TestCase + setup do + @user = users(:family_admin) + @family = @user.family + @budget = budgets(:one) + @function = Assistant::Function::UpdateBudget.new(@user) + end + + test "has correct name" do + assert_equal "update_budget", @function.name + end + + test "has a description" do + assert_not_empty @function.description + end + + test "is not in strict mode" do + refute @function.strict_mode? + end + + test "params_schema declares month totals and categories as optional" do + schema = @function.params_schema + assert schema[:properties].key?(:month) + assert schema[:properties].key?(:budgeted_spending) + assert schema[:properties].key?(:expected_income) + assert schema[:properties].key?(:categories) + assert_empty schema[:required] + end + + test "updates totals for the current month by default" do + result = @function.call( + "budgeted_spending" => 6500, + "expected_income" => 9000 + ) + + assert_equal true, result[:success] + assert_equal @budget.to_param, result[:month] + + @budget.reload + assert_equal 6500, @budget.budgeted_spending + assert_equal 9000, @budget.expected_income + end + + test "sets a category allocation by case-insensitive name" do + @budget.sync_budget_categories + + result = @function.call( + "categories" => [ { "category" => "food & drink", "amount" => 450 } ] + ) + + assert_equal true, result[:success] + assert_equal 1, result[:updated_categories].length + assert_equal "Food & Drink", result[:updated_categories].first[:category] + + budget_category = @budget.budget_categories + .joins(:category).find_by(categories: { name: "Food & Drink" }) + assert_equal 450, budget_category.reload.budgeted_spending + end + + test "sets a category allocation by id and keeps a subcategory's parent in sync" do + @budget.sync_budget_categories + subcategory = categories(:subcategory) + parent = subcategory.parent + + result = @function.call( + "categories" => [ { "category" => subcategory.id, "amount" => 120 } ] + ) + + assert_equal true, result[:success] + + sub_bc = @budget.budget_categories.find_by(category_id: subcategory.id) + parent_bc = @budget.budget_categories.find_by(category_id: parent.id) + assert_equal 120, sub_bc.reload.budgeted_spending + assert_operator parent_bc.reload.budgeted_spending, :>=, 120 + end + + test "bootstraps a budget when targeting a valid month with no budget row" do + target = Date.current.beginning_of_month << 1 + assert_nil @family.budgets.find_by(start_date: target) + + result = @function.call( + "month" => target.strftime("%Y-%m"), + "budgeted_spending" => 1000 + ) + + assert_equal true, result[:success] + created = @family.budgets.find_by(start_date: target) + assert_equal 1000, created.budgeted_spending + end + + test "rejects a category outside the family and rolls back totals from the same call" do + other_family_category = Category.create!( + family: families(:empty), + name: "Elsewhere", + color: "#e99537", + lucide_icon: "tag" + ) + original = @budget.budgeted_spending + + result = @function.call( + "budgeted_spending" => 9999, + "categories" => [ { "category" => other_family_category.id, "amount" => 10 } ] + ) + + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + assert_equal original, @budget.reload.budgeted_spending + end + + test "does not leave a bootstrapped budget behind when a later entry is invalid" do + target = Date.current.beginning_of_month << 2 + assert_nil @family.budgets.find_by(start_date: target) + + result = @function.call( + "month" => target.strftime("%Y-%m"), + "budgeted_spending" => 1000, + "categories" => [ { "category" => "No Such Category", "amount" => 10 } ] + ) + + assert_equal false, result[:success] + assert_nil @family.budgets.find_by(start_date: target) + end + + test "applies explicit parent amounts after subcategory updates regardless of order" do + @budget.sync_budget_categories + subcategory = categories(:subcategory) + parent = subcategory.parent + + result = @function.call( + "categories" => [ + { "category" => parent.id, "amount" => 1000 }, + { "category" => subcategory.id, "amount" => 300 } + ] + ) + + assert_equal true, result[:success] + + sub_bc = @budget.budget_categories.find_by(category_id: subcategory.id) + parent_bc = @budget.budget_categories.find_by(category_id: parent.id) + assert_equal 300, sub_bc.reload.budgeted_spending + assert_equal 1000, parent_bc.reload.budgeted_spending + end + + test "explains that Uncategorized cannot be set directly" do + @budget.sync_budget_categories + + result = @function.call( + "categories" => [ { "category" => "Uncategorized", "amount" => 50 } ] + ) + + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + assert_match(/cannot be set directly/i, result[:message]) + end + + test "rejects unknown category names" do + result = @function.call( + "categories" => [ { "category" => "No Such Category", "amount" => 10 } ] + ) + + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + assert_match(/not found/i, result[:message]) + end + + test "rejects negative amounts" do + result = @function.call("expected_income" => -5) + + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + end + + test "rejects non-finite amounts" do + @budget.sync_budget_categories + + [ Float::NAN, Float::INFINITY ].each do |value| + result = @function.call("budgeted_spending" => value) + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + + result = @function.call("categories" => [ { "category" => "Food & Drink", "amount" => value } ]) + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + end + end + + test "rejects malformed months" do + result = @function.call("month" => "never", "budgeted_spending" => 1) + + assert_equal false, result[:success] + assert_equal "invalid_params", result[:error] + end + + test "rejects months outside the valid budget range" do + far_future = (Date.current + 10.years).strftime("%Y-%m") + + result = @function.call("month" => far_future, "budgeted_spending" => 1) + + assert_equal false, result[:success] + assert_equal "invalid_month", result[:error] + end + + test "rejects calls with nothing to change" do + result = @function.call({}) + + assert_equal false, result[:success] + assert_equal "no_changes", result[:error] + end +end