zk-data-agent/tests/test_worktree_runtime.py

from __future__ import annotations

import json
import shutil
import subprocess
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from src.agent_runtime import LocalCodingAgent
from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
from src.agent_types import AgentPermissions, AgentRuntimeConfig, ModelConfig
from src.worktree_runtime import WorktreeRuntime


class _FakeHTTPResponse:
    def __init__(self, payload: dict[str, object]) -> None:
        self.payload = payload

    def read(self) -> bytes:
        return json.dumps(self.payload).encode('utf-8')

    def __enter__(self) -> '_FakeHTTPResponse':
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        return None


def _make_urlopen_side_effect(responses: list[dict[str, object]]):
    queued = [_FakeHTTPResponse(payload) for payload in responses]

    def _fake_urlopen(request_obj, timeout=None):  # noqa: ANN001
        return queued.pop(0)

    return _fake_urlopen


def _init_git_repo(workspace: Path) -> None:
    subprocess.run(['git', 'init', '-q'], cwd=workspace, check=True)
    subprocess.run(['git', 'config', 'user.email', 'test@example.com'], cwd=workspace, check=True)
    subprocess.run(['git', 'config', 'user.name', 'Test User'], cwd=workspace, check=True)
    (workspace / 'README.md').write_text('hello\n', encoding='utf-8')
    subprocess.run(['git', 'add', 'README.md'], cwd=workspace, check=True)
    subprocess.run(['git', 'commit', '-qm', 'init'], cwd=workspace, check=True)


@unittest.skipUnless(shutil.which('git'), 'git is required for worktree tests')
class WorktreeRuntimeTests(unittest.TestCase):
    def test_worktree_runtime_enters_and_exits_managed_session(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            workspace = Path(tmp_dir)
            _init_git_repo(workspace)
            runtime = WorktreeRuntime.from_workspace(workspace)
            enter_report = runtime.enter('feature-preview')
            worktree_path = Path(enter_report.worktree_path or '')
            exit_report = runtime.exit(action='keep')

        self.assertTrue(enter_report.active)
        self.assertTrue(worktree_path.exists())
        self.assertIn('feature-preview', enter_report.worktree_branch or '')
        self.assertFalse(exit_report.active)
        self.assertEqual(exit_report.original_cwd, str(workspace))

    def test_worktree_tools_execute_against_runtime(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            workspace = Path(tmp_dir)
            _init_git_repo(workspace)
            runtime = WorktreeRuntime.from_workspace(workspace)
            context = build_tool_context(
                AgentRuntimeConfig(
                    cwd=workspace,
                    permissions=AgentPermissions(allow_file_write=True),
                ),
                worktree_runtime=runtime,
            )
            enter_result = execute_tool(
                default_tool_registry(),
                'worktree_enter',
                {'name': 'preview'},
                context,
            )
            status_result = execute_tool(
                default_tool_registry(),
                'worktree_status',
                {},
                context,
            )
            exit_result = execute_tool(
                default_tool_registry(),
                'worktree_exit',
                {'action': 'remove', 'discard_changes': True},
                context,
            )

        self.assertTrue(enter_result.ok)
        self.assertIn('preview', enter_result.content)
        self.assertEqual(enter_result.metadata.get('action'), 'worktree_enter')
        self.assertTrue(status_result.ok)
        self.assertIn('Active managed worktree: True', status_result.content)
        self.assertTrue(exit_result.ok)
        self.assertEqual(exit_result.metadata.get('action'), 'worktree_exit')

    def test_agent_switches_cwd_after_worktree_enter(self) -> None:
        responses = [
            {
                'choices': [
                    {
                        'message': {
                            'role': 'assistant',
                            'content': 'Entering worktree.',
                            'tool_calls': [
                                {
                                    'id': 'call_enter',
                                    'type': 'function',
                                    'function': {
                                        'name': 'worktree_enter',
                                        'arguments': '{"name":"preview"}',
                                    },
                                }
                            ],
                        },
                        'finish_reason': 'tool_calls',
                    }
                ],
                'usage': {'prompt_tokens': 8, 'completion_tokens': 2},
            },
            {
                'choices': [
                    {
                        'message': {
                            'role': 'assistant',
                            'content': 'Writing inside the worktree.',
                            'tool_calls': [
                                {
                                    'id': 'call_write',
                                    'type': 'function',
                                    'function': {
                                        'name': 'write_file',
                                        'arguments': '{"path":"note.txt","content":"from worktree\\n"}',
                                    },
                                }
                            ],
                        },
                        'finish_reason': 'tool_calls',
                    }
                ],
                'usage': {'prompt_tokens': 8, 'completion_tokens': 2},
            },
            {
                'choices': [
                    {
                        'message': {
                            'role': 'assistant',
                            'content': 'done',
                        },
                        'finish_reason': 'stop',
                    }
                ],
                'usage': {'prompt_tokens': 6, 'completion_tokens': 1},
            },
        ]
        with tempfile.TemporaryDirectory() as tmp_dir:
            workspace = Path(tmp_dir)
            _init_git_repo(workspace)
            with patch(
                'src.openai_compat.request.urlopen',
                side_effect=_make_urlopen_side_effect(responses),
            ):
                agent = LocalCodingAgent(
                    model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
                    runtime_config=AgentRuntimeConfig(
                        cwd=workspace,
                        permissions=AgentPermissions(allow_file_write=True),
                    ),
                )
                result = agent.run('Use a worktree and write a file there')
                runtime = WorktreeRuntime.from_workspace(workspace)
                assert runtime.active_session is not None
                worktree_path = Path(runtime.active_session.worktree_path)

        self.assertEqual(result.final_output, 'done')
        self.assertFalse((workspace / 'note.txt').exists())
        self.assertTrue((worktree_path / 'note.txt').exists())
        self.assertEqual(agent.runtime_config.cwd, worktree_path.resolve())