Add tests to ensure run limits are respected

src/controlflow/agents/agent.py

@@ -282,6 +282,6 @@ async def _run_model_async(
         yield AgentMessage(agent=self, message=response)
 
         for tool_call in response.tool_calls + response.invalid_tool_calls:
-            yield ToolCallEvent(agent=self, tool_call=tool_call, message=response)
+            yield ToolCallEvent(agent=self, tool_call=tool_call)
             result = await handle_tool_call_async(tool_call, tools=tools)
             yield ToolResultEvent(agent=self, tool_call=tool_call, tool_result=result)

src/controlflow/tasks/task.py

@@ -677,20 +677,28 @@ def run(
     objective: str,
     *task_args,
     turn_strategy: "TurnStrategy" = None,
+    max_calls_per_turn: int = None,
+    max_turns: int = None,
     **task_kwargs,
 ):
     task = controlflow.Task(
         objective=objective,
         *task_args,
         **task_kwargs,
     )
-    return task.run(turn_strategy=turn_strategy)
+    return task.run(
+        turn_strategy=turn_strategy,
+        max_calls_per_turn=max_calls_per_turn,
+        max_turns=max_turns,
+    )
 
 
 async def run_async(
     objective: str,
     *task_args,
     turn_strategy: "TurnStrategy" = None,
+    max_calls_per_turn: int = None,
+    max_turns: int = None,
     **task_kwargs,
 ):
     task = controlflow.Task(
@@ -701,4 +709,6 @@ async def run_async(
 
     return await task.run_async(
         turn_strategy=turn_strategy,
+        max_calls_per_turn=max_calls_per_turn,
+        max_turns=max_turns,
     )

tests/tasks/test_tasks.py

@@ -331,3 +331,37 @@ def test_success_tool_with_list_of_options_requires_int(self):
         tool = task.create_success_tool()
         with pytest.raises(ValueError):
             tool.run(input=dict(result="good"))
+
+
+@pytest.mark.parametrize(
+    "max_turns, max_calls_per_turn, expected_calls",
+    [
+        (1, 1, 1),
+        (1, 2, 2),
+        (2, 1, 2),
+        (3, 2, 6),
+    ],
+)
+def test_run_with_limits(
+    monkeypatch, default_fake_llm, max_turns, max_calls_per_turn, expected_calls
+):
+    # Tests that the run function correctly limits the number of turns and calls per turn
+    default_fake_llm.set_responses(["hello", "world", "how", "are", "you"])
+
+    call_count = 0
+    original_run_model = Agent._run_model
+
+    def mock_run_model(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        return original_run_model(*args, **kwargs)
+
+    monkeypatch.setattr(Agent, "_run_model", mock_run_model)
+
+    task = Task("send messages")
+    task.run(
+        max_calls_per_turn=max_calls_per_turn,
+        max_turns=max_turns,
+    )
+
+    assert call_count == expected_calls

tests/test_run.py

@@ -0,0 +1,48 @@
+import pytest
+
+import controlflow
+from controlflow.agents.agent import Agent
+
+
+def test_run():
+    result = controlflow.run("what's 2 + 2", result_type=int)
+    assert result == 4
+
+
+async def test_run_async():
+    result = await controlflow.run_async("what's 2 + 2", result_type=int)
+    assert result == 4
+
+
+@pytest.mark.parametrize(
+    "max_turns, max_calls_per_turn, expected_calls",
+    [
+        (1, 1, 1),
+        (1, 2, 2),
+        (2, 1, 2),
+        (3, 2, 6),
+    ],
+)
+def test_run_with_limits(
+    monkeypatch, default_fake_llm, max_turns, max_calls_per_turn, expected_calls
+):
+    # Tests that the run function correctly limits the number of turns and calls per turn
+    default_fake_llm.set_responses(["hello", "world", "how", "are", "you"])
+
+    call_count = 0
+    original_run_model = Agent._run_model
+
+    def mock_run_model(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        return original_run_model(*args, **kwargs)
+
+    monkeypatch.setattr(Agent, "_run_model", mock_run_model)
+
+    controlflow.run(
+        "send messages",
+        max_calls_per_turn=max_calls_per_turn,
+        max_turns=max_turns,
+    )
+
+    assert call_count == expected_calls