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:
Zamil Majdy
2026-03-27 10:04:13 +07:00
parent 3d8c040e33
commit 96554732ee
5 changed files with 140 additions and 241 deletions

View File

@@ -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)

View File

@@ -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()

View File

@@ -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.

View File

@@ -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({})

View File

@@ -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()
)
# ---------------------------------------------------------------------------