mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-04-08 03:00:28 -04:00
refactor(copilot): remove _session_dry_run ContextVar; use session.dry_run directly
session.dry_run (via ChatSession.metadata) is the single source of truth. Remove the redundant _session_dry_run ContextVar, is_session_dry_run() helper, and dry_run parameter from set_execution_context(). Tool handlers now read session.dry_run directly.
This commit is contained in:
@@ -204,9 +204,9 @@ async def stream_chat_completion_baseline(
|
||||
|
||||
tools = get_available_tools()
|
||||
|
||||
# Propagate execution context (user, session, dry_run) so tool handlers
|
||||
# (run_block, run_agent) can read session-level flags via ContextVars.
|
||||
set_execution_context(user_id, session, dry_run=session.dry_run)
|
||||
# Propagate execution context so tool handlers can read session-level
|
||||
# flags. dry_run is read directly from session.dry_run by each tool.
|
||||
set_execution_context(user_id, session)
|
||||
|
||||
yield StreamStart(messageId=message_id, sessionId=session_id)
|
||||
|
||||
|
||||
@@ -52,11 +52,6 @@ _current_permissions: "ContextVar[CopilotPermissions | None]" = ContextVar(
|
||||
"_current_permissions", default=None
|
||||
)
|
||||
|
||||
# Session-level dry-run flag. When True, all tool calls (run_block,
|
||||
# run_agent) are forced to use dry-run simulation — no real API calls or
|
||||
# side effects. Set by set_execution_context(); read by tool handlers.
|
||||
_session_dry_run: ContextVar[bool] = ContextVar("_session_dry_run", default=False)
|
||||
|
||||
|
||||
def encode_cwd_for_cli(cwd: str) -> str:
|
||||
"""Encode a working directory path the same way the Claude CLI does.
|
||||
@@ -78,7 +73,6 @@ def set_execution_context(
|
||||
sandbox: "AsyncSandbox | None" = None,
|
||||
sdk_cwd: str | None = None,
|
||||
permissions: "CopilotPermissions | None" = None,
|
||||
dry_run: bool = False,
|
||||
) -> None:
|
||||
"""Set per-turn context variables used by file-resolution tool handlers."""
|
||||
_current_user_id.set(user_id)
|
||||
@@ -87,7 +81,6 @@ def set_execution_context(
|
||||
_current_sdk_cwd.set(sdk_cwd or "")
|
||||
_current_project_dir.set(_encode_cwd_for_cli(sdk_cwd) if sdk_cwd else "")
|
||||
_current_permissions.set(permissions)
|
||||
_session_dry_run.set(dry_run)
|
||||
|
||||
|
||||
def get_execution_context() -> tuple[str | None, ChatSession | None]:
|
||||
@@ -100,11 +93,6 @@ def get_current_permissions() -> "CopilotPermissions | None":
|
||||
return _current_permissions.get()
|
||||
|
||||
|
||||
def is_session_dry_run() -> bool:
|
||||
"""Return True if the current session has session-level dry-run enabled."""
|
||||
return _session_dry_run.get()
|
||||
|
||||
|
||||
def get_current_sandbox() -> "AsyncSandbox | None":
|
||||
"""Return the E2B sandbox for the current session, or None if not active."""
|
||||
return _current_sandbox.get()
|
||||
|
||||
@@ -1772,7 +1772,6 @@ async def stream_chat_completion_sdk(
|
||||
sandbox=e2b_sandbox,
|
||||
sdk_cwd=sdk_cwd,
|
||||
permissions=permissions,
|
||||
dry_run=session.dry_run,
|
||||
)
|
||||
|
||||
# Fail fast when no API credentials are available at all.
|
||||
|
||||
@@ -23,7 +23,6 @@ from backend.copilot.context import (
|
||||
_current_session,
|
||||
_current_user_id,
|
||||
_encode_cwd_for_cli,
|
||||
_session_dry_run,
|
||||
get_execution_context,
|
||||
get_sdk_cwd,
|
||||
is_allowed_local_path,
|
||||
@@ -95,7 +94,6 @@ def set_execution_context(
|
||||
sandbox: "AsyncSandbox | None" = None,
|
||||
sdk_cwd: str | None = None,
|
||||
permissions: "CopilotPermissions | None" = None,
|
||||
dry_run: bool = False,
|
||||
) -> None:
|
||||
"""Set the execution context for tool calls.
|
||||
|
||||
@@ -108,8 +106,6 @@ def set_execution_context(
|
||||
sandbox: Optional E2B sandbox; when set, bash_exec routes commands there.
|
||||
sdk_cwd: SDK working directory; used to scope tool-results reads.
|
||||
permissions: Optional capability filter restricting tools/blocks.
|
||||
dry_run: When True, session-level dry-run is active — all tool calls
|
||||
(run_block, run_agent) are forced to use dry-run simulation.
|
||||
"""
|
||||
_current_user_id.set(user_id)
|
||||
_current_session.set(session)
|
||||
@@ -117,7 +113,6 @@ def set_execution_context(
|
||||
_current_sdk_cwd.set(sdk_cwd or "")
|
||||
_current_project_dir.set(_encode_cwd_for_cli(sdk_cwd) if sdk_cwd else "")
|
||||
_current_permissions.set(permissions)
|
||||
_session_dry_run.set(dry_run)
|
||||
_pending_tool_outputs.set({})
|
||||
_stash_event.set(asyncio.Event())
|
||||
_tool_task_queues.set({})
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
Verifies that when a session has dry_run=True, all tool calls (run_block,
|
||||
run_agent) are forced to use dry-run mode, regardless of what the individual
|
||||
tool call specifies.
|
||||
tool call specifies. The single source of truth is ``session.dry_run``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,11 +11,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from backend.copilot.context import (
|
||||
_session_dry_run,
|
||||
is_session_dry_run,
|
||||
set_execution_context,
|
||||
)
|
||||
from backend.copilot.model import ChatSession
|
||||
from backend.copilot.tools.models import ErrorResponse
|
||||
from backend.copilot.tools.run_agent import RunAgentInput, RunAgentTool
|
||||
@@ -60,61 +55,6 @@ def _make_mock_block(name: str = "TestBlock"):
|
||||
return block
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ContextVar tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSessionDryRunContextVar:
|
||||
"""Test the session dry_run ContextVar."""
|
||||
|
||||
def test_default_is_false(self):
|
||||
"""Default value of _session_dry_run should be False."""
|
||||
# Reset to default
|
||||
token = _session_dry_run.set(False)
|
||||
try:
|
||||
assert is_session_dry_run() is False
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
|
||||
def test_set_true(self):
|
||||
"""Setting _session_dry_run to True should be reflected by is_session_dry_run."""
|
||||
token = _session_dry_run.set(True)
|
||||
try:
|
||||
assert is_session_dry_run() is True
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
|
||||
|
||||
class TestSetExecutionContextDryRun:
|
||||
"""Test that set_execution_context propagates dry_run."""
|
||||
|
||||
def test_sets_dry_run_true(self):
|
||||
token = _session_dry_run.set(False)
|
||||
try:
|
||||
session = _make_session(dry_run=True)
|
||||
set_execution_context(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
dry_run=True,
|
||||
)
|
||||
assert is_session_dry_run() is True
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
|
||||
def test_sets_dry_run_false_by_default(self):
|
||||
token = _session_dry_run.set(True)
|
||||
try:
|
||||
session = _make_session(dry_run=False)
|
||||
set_execution_context(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
)
|
||||
assert is_session_dry_run() is False
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RunBlockTool tests
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -126,106 +66,95 @@ class TestRunBlockToolSessionDryRun:
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_dry_run_forces_block_dry_run(self):
|
||||
"""When session dry_run is True, run_block should force dry_run=True."""
|
||||
token = _session_dry_run.set(True)
|
||||
try:
|
||||
tool = RunBlockTool()
|
||||
session = _make_session(dry_run=True)
|
||||
tool = RunBlockTool()
|
||||
session = _make_session(dry_run=True)
|
||||
|
||||
mock_block = _make_mock_block()
|
||||
mock_block = _make_mock_block()
|
||||
|
||||
async def fake_simulate(block, input_data):
|
||||
yield "result", "simulated"
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.prepare_block_for_execution"
|
||||
) as mock_prep,
|
||||
patch("backend.copilot.tools.run_block.execute_block") as mock_exec,
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.get_current_permissions",
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
# Set up prepare_block_for_execution to return a mock prep
|
||||
mock_prep_result = MagicMock()
|
||||
mock_prep_result.block = mock_block
|
||||
mock_prep_result.input_data = {"query": "test"}
|
||||
mock_prep_result.matched_credentials = {}
|
||||
mock_prep_result.synthetic_node_id = "node-1"
|
||||
mock_prep.return_value = mock_prep_result
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.prepare_block_for_execution"
|
||||
) as mock_prep,
|
||||
patch("backend.copilot.tools.run_block.execute_block") as mock_exec,
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.get_current_permissions",
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
# Set up prepare_block_for_execution to return a mock prep
|
||||
mock_prep_result = MagicMock()
|
||||
mock_prep_result.block = mock_block
|
||||
mock_prep_result.input_data = {"query": "test"}
|
||||
mock_prep_result.matched_credentials = {}
|
||||
mock_prep_result.synthetic_node_id = "node-1"
|
||||
mock_prep.return_value = mock_prep_result
|
||||
# Set up execute_block to return a success
|
||||
mock_exec.return_value = MagicMock(
|
||||
message="[DRY RUN] Block executed",
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Set up execute_block to return a success
|
||||
mock_exec.return_value = MagicMock(
|
||||
message="[DRY RUN] Block executed",
|
||||
success=True,
|
||||
)
|
||||
await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
block_id="test-block-id",
|
||||
input_data={"query": "test"},
|
||||
dry_run=False, # User passed False, but session overrides
|
||||
)
|
||||
|
||||
await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
block_id="test-block-id",
|
||||
input_data={"query": "test"},
|
||||
dry_run=False, # User passed False, but session overrides
|
||||
)
|
||||
|
||||
# Verify execute_block was called with dry_run=True
|
||||
mock_exec.assert_called_once()
|
||||
call_kwargs = mock_exec.call_args
|
||||
assert call_kwargs.kwargs.get("dry_run") is True
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
# Verify execute_block was called with dry_run=True
|
||||
mock_exec.assert_called_once()
|
||||
call_kwargs = mock_exec.call_args
|
||||
assert call_kwargs.kwargs.get("dry_run") is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_session_dry_run_respects_tool_param(self):
|
||||
"""When session dry_run is False, tool-level dry_run should be respected."""
|
||||
token = _session_dry_run.set(False)
|
||||
try:
|
||||
tool = RunBlockTool()
|
||||
session = _make_session(dry_run=False)
|
||||
tool = RunBlockTool()
|
||||
session = _make_session(dry_run=False)
|
||||
|
||||
mock_block = _make_mock_block()
|
||||
mock_block = _make_mock_block()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.prepare_block_for_execution"
|
||||
) as mock_prep,
|
||||
patch("backend.copilot.tools.run_block.execute_block") as mock_exec,
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.get_current_permissions",
|
||||
return_value=None,
|
||||
),
|
||||
patch("backend.copilot.tools.run_block.check_hitl_review") as mock_hitl,
|
||||
):
|
||||
mock_prep_result = MagicMock()
|
||||
mock_prep_result.block = mock_block
|
||||
mock_prep_result.input_data = {"query": "test"}
|
||||
mock_prep_result.matched_credentials = {}
|
||||
mock_prep_result.synthetic_node_id = "node-1"
|
||||
mock_prep_result.required_non_credential_keys = {"query"}
|
||||
mock_prep_result.provided_input_keys = {"query"}
|
||||
mock_prep.return_value = mock_prep_result
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.prepare_block_for_execution"
|
||||
) as mock_prep,
|
||||
patch("backend.copilot.tools.run_block.execute_block") as mock_exec,
|
||||
patch(
|
||||
"backend.copilot.tools.run_block.get_current_permissions",
|
||||
return_value=None,
|
||||
),
|
||||
patch("backend.copilot.tools.run_block.check_hitl_review") as mock_hitl,
|
||||
):
|
||||
mock_prep_result = MagicMock()
|
||||
mock_prep_result.block = mock_block
|
||||
mock_prep_result.input_data = {"query": "test"}
|
||||
mock_prep_result.matched_credentials = {}
|
||||
mock_prep_result.synthetic_node_id = "node-1"
|
||||
mock_prep_result.required_non_credential_keys = {"query"}
|
||||
mock_prep_result.provided_input_keys = {"query"}
|
||||
mock_prep.return_value = mock_prep_result
|
||||
|
||||
mock_hitl.return_value = ("node-exec-1", {"query": "test"})
|
||||
mock_hitl.return_value = ("node-exec-1", {"query": "test"})
|
||||
|
||||
mock_exec.return_value = MagicMock(
|
||||
message="Block executed",
|
||||
success=True,
|
||||
)
|
||||
mock_exec.return_value = MagicMock(
|
||||
message="Block executed",
|
||||
success=True,
|
||||
)
|
||||
|
||||
await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
block_id="test-block-id",
|
||||
input_data={"query": "test"},
|
||||
dry_run=False,
|
||||
)
|
||||
await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
block_id="test-block-id",
|
||||
input_data={"query": "test"},
|
||||
dry_run=False,
|
||||
)
|
||||
|
||||
# Verify execute_block was called with dry_run=False
|
||||
mock_exec.assert_called_once()
|
||||
call_kwargs = mock_exec.call_args
|
||||
assert call_kwargs.kwargs.get("dry_run") is False
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
# Verify execute_block was called with dry_run=False
|
||||
mock_exec.assert_called_once()
|
||||
call_kwargs = mock_exec.call_args
|
||||
assert call_kwargs.kwargs.get("dry_run") is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -239,92 +168,80 @@ class TestRunAgentToolSessionDryRun:
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_dry_run_forces_agent_dry_run(self):
|
||||
"""When session dry_run is True, run_agent params.dry_run should be forced True."""
|
||||
token = _session_dry_run.set(True)
|
||||
try:
|
||||
tool = RunAgentTool()
|
||||
session = _make_session(dry_run=True)
|
||||
tool = RunAgentTool()
|
||||
session = _make_session(dry_run=True)
|
||||
|
||||
# Mock the graph and dependencies
|
||||
mock_graph = MagicMock()
|
||||
mock_graph.id = "graph-1"
|
||||
mock_graph.name = "Test Agent"
|
||||
mock_graph.description = "A test agent"
|
||||
mock_graph.input_schema = {"properties": {}, "required": []}
|
||||
mock_graph.trigger_setup_info = None
|
||||
# Mock the graph and dependencies
|
||||
mock_graph = MagicMock()
|
||||
mock_graph.id = "graph-1"
|
||||
mock_graph.name = "Test Agent"
|
||||
mock_graph.description = "A test agent"
|
||||
mock_graph.input_schema = {"properties": {}, "required": []}
|
||||
mock_graph.trigger_setup_info = None
|
||||
|
||||
mock_library_agent = MagicMock()
|
||||
mock_library_agent.id = "lib-1"
|
||||
mock_library_agent.graph_id = "graph-1"
|
||||
mock_library_agent.graph_version = 1
|
||||
mock_library_agent.name = "Test Agent"
|
||||
mock_library_agent = MagicMock()
|
||||
mock_library_agent.id = "lib-1"
|
||||
mock_library_agent.graph_id = "graph-1"
|
||||
mock_library_agent.graph_version = 1
|
||||
mock_library_agent.name = "Test Agent"
|
||||
|
||||
mock_execution = MagicMock()
|
||||
mock_execution.id = "exec-1"
|
||||
mock_execution = MagicMock()
|
||||
mock_execution.id = "exec-1"
|
||||
|
||||
with (
|
||||
patch("backend.copilot.tools.run_agent.graph_db"),
|
||||
patch("backend.copilot.tools.run_agent.library_db"),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.fetch_graph_from_store_slug",
|
||||
return_value=(mock_graph, None),
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.match_user_credentials_to_graph",
|
||||
return_value=({}, []),
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.get_or_create_library_agent",
|
||||
return_value=mock_library_agent,
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.execution_utils"
|
||||
) as mock_exec_utils,
|
||||
patch("backend.copilot.tools.run_agent.track_agent_run_success"),
|
||||
):
|
||||
mock_exec_utils.add_graph_execution = AsyncMock(
|
||||
return_value=mock_execution
|
||||
)
|
||||
with (
|
||||
patch("backend.copilot.tools.run_agent.graph_db"),
|
||||
patch("backend.copilot.tools.run_agent.library_db"),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.fetch_graph_from_store_slug",
|
||||
return_value=(mock_graph, None),
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.match_user_credentials_to_graph",
|
||||
return_value=({}, []),
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.get_or_create_library_agent",
|
||||
return_value=mock_library_agent,
|
||||
),
|
||||
patch("backend.copilot.tools.run_agent.execution_utils") as mock_exec_utils,
|
||||
patch("backend.copilot.tools.run_agent.track_agent_run_success"),
|
||||
):
|
||||
mock_exec_utils.add_graph_execution = AsyncMock(return_value=mock_execution)
|
||||
|
||||
await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
username_agent_slug="user/test-agent",
|
||||
dry_run=False, # User passed False, but session overrides
|
||||
use_defaults=True,
|
||||
)
|
||||
await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
username_agent_slug="user/test-agent",
|
||||
dry_run=False, # User passed False, but session overrides
|
||||
use_defaults=True,
|
||||
)
|
||||
|
||||
# Verify add_graph_execution was called with dry_run=True
|
||||
mock_exec_utils.add_graph_execution.assert_called_once()
|
||||
call_kwargs = mock_exec_utils.add_graph_execution.call_args
|
||||
assert call_kwargs.kwargs.get("dry_run") is True
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
# Verify add_graph_execution was called with dry_run=True
|
||||
mock_exec_utils.add_graph_execution.assert_called_once()
|
||||
call_kwargs = mock_exec_utils.add_graph_execution.call_args
|
||||
assert call_kwargs.kwargs.get("dry_run") is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_dry_run_blocks_scheduling(self):
|
||||
"""When session dry_run is True, scheduling requests should be rejected."""
|
||||
token = _session_dry_run.set(True)
|
||||
try:
|
||||
tool = RunAgentTool()
|
||||
session = _make_session(dry_run=True)
|
||||
tool = RunAgentTool()
|
||||
session = _make_session(dry_run=True)
|
||||
|
||||
result = await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
username_agent_slug="user/test-agent",
|
||||
schedule_name="daily-run",
|
||||
cron="0 9 * * *",
|
||||
dry_run=False, # Session overrides to True
|
||||
)
|
||||
result = await tool._execute(
|
||||
user_id="test-user",
|
||||
session=session,
|
||||
username_agent_slug="user/test-agent",
|
||||
schedule_name="daily-run",
|
||||
cron="0 9 * * *",
|
||||
dry_run=False, # Session overrides to True
|
||||
)
|
||||
|
||||
assert isinstance(result, ErrorResponse)
|
||||
assert "dry-run" in result.message.lower()
|
||||
assert (
|
||||
"scheduling" in result.message.lower()
|
||||
or "schedule" in result.message.lower()
|
||||
)
|
||||
finally:
|
||||
_session_dry_run.reset(token)
|
||||
assert isinstance(result, ErrorResponse)
|
||||
assert "dry-run" in result.message.lower()
|
||||
assert (
|
||||
"scheduling" in result.message.lower()
|
||||
or "schedule" in result.message.lower()
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user