diff --git a/third_party/blink/renderer/core/script_tools/model_context.cc b/third_party/blink/renderer/core/script_tools/model_context.cc index 0fe934da99569..25cbb58fbb658 100644 --- a/third_party/blink/renderer/core/script_tools/model_context.cc +++ b/third_party/blink/renderer/core/script_tools/model_context.cc @@ -695,20 +695,22 @@ void ModelContext::RegisterDeclarativeTool( void ModelContext::OnToolExecuted( const base::UnguessableToken& invocation_id, base::expected> result) { - auto it = pending_executions_.find(String(invocation_id.ToString())); - if (it == pending_executions_.end()) { + // Remove the entry up front because the probes and the completion callback + // below may re-enter this object and mutate `pending_executions_`. + PendingExecution pending_execution = + pending_executions_.Take(String(invocation_id.ToString())); + if (!pending_execution.callback) { return; } if (result.has_value()) { probe::WebMCPToolResponded(document_, result.value(), invocation_id); - std::move(it->value.callback).Run(result.value()); + std::move(pending_execution.callback).Run(result.value()); } else { ScriptToolError error(ScriptToolErrorCode::kToolInvocationFailed); probe::WebMCPToolFailed(document_, error, invocation_id, result.error()); - std::move(it->value.callback).Run(base::unexpected(error)); + std::move(pending_execution.callback).Run(base::unexpected(error)); } - pending_executions_.erase(it); } void ModelContext::MaybeRecordToolCount() { diff --git a/third_party/blink/renderer/core/script_tools/model_context_test.cc b/third_party/blink/renderer/core/script_tools/model_context_test.cc index ca2a149ce735b..df417eab006f7 100644 --- a/third_party/blink/renderer/core/script_tools/model_context_test.cc +++ b/third_party/blink/renderer/core/script_tools/model_context_test.cc @@ -1376,6 +1376,59 @@ TEST_F(ModelContextTest, CancelToolReentrancy) { run_loop.Run(); } +TEST_F(ModelContextTest, OnToolExecutedReentrancy) { + SimRequest main_resource("https://example.com/", "text/html"); + LoadURL("https://example.com/"); + + main_resource.Complete(R"( + + + )"); + + auto* model_context = ModelContextSupplement::modelContext(GetDocument()); + ASSERT_TRUE(model_context); + + // Start enough long-running invocations that cancelling them all will cause + // the pending executions map to shrink its backing store. + Vector hang_ids; + for (int i = 0; i < 32; ++i) { + base::UnguessableToken id = base::UnguessableToken::Create(); + hang_ids.push_back(id); + model_context->ExecuteTool(id, "hang", "{}", /* signal= */ nullptr, + base::DoNothing()); + } + + base::RunLoop run_loop; + String result; + model_context->ExecuteTool( + base::UnguessableToken::Create(), "echo", "{}", /* signal= */ nullptr, + base::BindLambdaForTesting( + [&](base::expected res) { + // Cancel all other pending invocations from inside the completion + // callback. + for (const auto& id : hang_ids) { + model_context->CancelTool(id); + } + ASSERT_TRUE(res.has_value()); + result = *res; + run_loop.Quit(); + })); + + run_loop.Run(); + EXPECT_EQ(result, "done"); +} + class MockDeclarativeTool : public GarbageCollected, public DeclarativeWebMCPTool { public: