Merge branch 'develop' into eagarwal-cg-10796-lsp-progress-reporting-support

Author: bagel897

.github/actions/run-ats/action.yml

@@ -48,11 +48,10 @@ runs:
           echo "No tests to run, skipping..."
           exit 0
         fi
-        echo $TESTS_TO_RUN | xargs uv run pytest --cov \
+        echo $TESTS_TO_RUN | xargs uv run --frozen pytest --cov \
           -o junit_suite_name="${{ github.job }}" \
           -n auto \
           -vv \
-          --cov \
           --cov-append \
           ${{ inputs.collect_args }}
 

src/codegen/extensions/lsp/io.py

@@ -16,7 +16,7 @@
 
 @dataclass
 class File:
-    doc: TextDocument
+    doc: TextDocument | None
     path: Path
     change: TextEdit | None = None
     other_change: CreateFile | RenameFile | DeleteFile | None = None
@@ -65,6 +65,8 @@ def read_text(self, path: Path) -> str:
             return file.change.new_text
         if file.created:
             return ""
+        if file.doc is None:
+            return self.base_io.read_text(path)
         return file.doc.source
 
     def read_bytes(self, path: Path) -> bytes:
@@ -76,6 +78,8 @@ def read_bytes(self, path: Path) -> bytes:
             return file.change.new_text.encode("utf-8")
         if file.created:
             return b""
+        if file.doc is None:
+            return self.base_io.read_bytes(path)
         return file.doc.source.encode("utf-8")
 
     def write_bytes(self, path: Path, content: bytes) -> None:
@@ -112,6 +116,8 @@ def file_exists(self, path: Path) -> bool:
             return True
         if file.created:
             return True
+        if file.doc is None:
+            return self.base_io.file_exists(path)
         try:
             file.doc.source
             return True
@@ -134,3 +140,13 @@ def get_workspace_edit(self) -> types.WorkspaceEdit:
                 file.change = None
         logger.info(f"Workspace edit: {pprint.pformat(list(map(asdict, document_changes)))}")
         return types.WorkspaceEdit(document_changes=document_changes)
+
+    def update_file(self, path: Path, version: int | None = None) -> None:
+        file = self._get_file(path)
+        file.doc = self.workspace.get_text_document(path.as_uri())
+        if version is not None:
+            file.version = version
+
+    def close_file(self, path: Path) -> None:
+        file = self._get_file(path)
+        file.doc = None

src/codegen/extensions/lsp/lsp.py

@@ -24,6 +24,7 @@ def did_open(server: CodegenLanguageServer, params: types.DidOpenTextDocumentPar
     # The document is automatically added to the workspace by pygls
     # We can perform any additional processing here if needed
     path = get_path(params.text_document.uri)
+    server.io.update_file(path, params.text_document.version)
     file = server.codebase.get_file(str(path), optional=True)
     if not isinstance(file, SourceFile) and path.suffix in server.codebase.ctx.extensions:
         sync = DiffLite(change_type=ChangeType.Added, path=path)
@@ -37,6 +38,7 @@ def did_change(server: CodegenLanguageServer, params: types.DidChangeTextDocumen
     # The document is automatically updated in the workspace by pygls
     # We can perform any additional processing here if needed
     path = get_path(params.text_document.uri)
+    server.io.update_file(path, params.text_document.version)
     sync = DiffLite(change_type=ChangeType.Modified, path=path)
     server.codebase.ctx.apply_diffs([sync])
 
@@ -63,6 +65,8 @@ def did_close(server: CodegenLanguageServer, params: types.DidCloseTextDocumentP
     logger.info(f"Document closed: {params.text_document.uri}")
     # The document is automatically removed from the workspace by pygls
     # We can perform any additional cleanup here if needed
+    path = get_path(params.text_document.uri)
+    server.io.close_file(path)
 
 
 @server.feature(

src/codegen/git/repo_operator/repo_operator.py

@@ -19,6 +19,7 @@
 from codegen.git.schemas.enums import CheckoutResult, FetchResult
 from codegen.git.schemas.repo_config import RepoConfig
 from codegen.git.utils.remote_progress import CustomRemoteProgress
+from codegen.shared.configs.session_configs import config
 from codegen.shared.performance.stopwatch_utils import stopwatch
 from codegen.shared.performance.time_utils import humanize_duration
 
@@ -46,7 +47,7 @@ def __init__(
     ) -> None:
         assert repo_config is not None
         self.repo_config = repo_config
-        self.access_token = access_token
+        self.access_token = access_token or config.secrets.github_token
         self.base_dir = repo_config.base_dir
         self.bot_commit = bot_commit
 

src/codegen/sdk/core/codebase.py

@@ -514,7 +514,7 @@ def get_file_from_path(path: Path) -> File | None:
         if file is not None:
             return file
         absolute_path = self.ctx.to_absolute(filepath)
-        if absolute_path.suffix in self.ctx.extensions:
+        if absolute_path.suffix in self.ctx.extensions and not self.ctx.io.file_exists(absolute_path):
             return None
         if self.ctx.io.file_exists(absolute_path):
             return get_file_from_path(absolute_path)
@@ -902,10 +902,30 @@ def restore_stashed_changes(self):
     ####################################################################################################################
 
     def create_pr(self, title: str, body: str) -> PullRequest:
-        """Creates a PR from the current branch."""
+        """Creates a pull request from the current branch to the repository's default branch.
+
+        This method will:
+        1. Stage and commit any pending changes with the PR title as the commit message
+        2. Push the current branch to the remote repository
+        3. Create a pull request targeting the default branch
+
+        Args:
+            title (str): The title for the pull request
+            body (str): The description/body text for the pull request
+
+        Returns:
+            PullRequest: The created GitHub pull request object
+
+        Raises:
+            ValueError: If attempting to create a PR while in a detached HEAD state
+            ValueError: If the current branch is the default branch
+        """
         if self._op.git_cli.head.is_detached:
             msg = "Cannot make a PR from a detached HEAD"
             raise ValueError(msg)
+        if self._op.git_cli.active_branch.name == self._op.default_branch:
+            msg = "Cannot make a PR from the default branch"
+            raise ValueError(msg)
         self._op.stage_and_commit_all_changes(message=title)
         self._op.push_changes()
         return self._op.remote_git_repo.create_pull(head_branch_name=self._op.git_cli.active_branch.name, base_branch_name=self._op.default_branch, title=title, body=body)

tests/integration/codegen/git/codebase/conftest.py

@@ -25,7 +25,7 @@ def repo_config(tmpdir):
 def op(repo_config):
     os.chdir(repo_config.base_dir)
     GitRepo.clone_from(url=get_authenticated_clone_url_for_repo_config(repo_config, token=config.secrets.github_token), to_path=os.path.join(repo_config.base_dir, repo_config.name), depth=1)
-    op = LocalRepoOperator(repo_config=repo_config, access_token=config.secrets.github_token)
+    op = LocalRepoOperator(repo_config=repo_config)
     yield op
 
 

tests/integration/codegen/git/codebase/test_codebase_create_pr.py

@@ -8,7 +8,8 @@
 def test_codebase_create_pr_active_branch(codebase: Codebase):
     head = f"test-create-pr-{uuid.uuid4()}"
     codebase.checkout(branch=head, create_if_missing=True)
-    codebase.files[0].remove()
+    file = codebase.files[0]
+    file.remove()
     codebase.commit()
     pr = codebase.create_pr(title="test-create-pr title", body="test-create-pr body")
     assert pr.title == "test-create-pr title"
@@ -17,10 +18,21 @@ def test_codebase_create_pr_active_branch(codebase: Codebase):
     assert pr.state == "open"
     assert pr.head.ref == head
     assert pr.base.ref == "main"
+    assert pr.get_files().totalCount == 1
+    assert pr.get_files()[0].filename == file.file_path
 
 
 def test_codebase_create_pr_detached_head(codebase: Codebase):
     codebase.checkout(commit=codebase._op.git_cli.head.commit)  # move to detached head state
     with pytest.raises(ValueError) as exc_info:
         codebase.create_pr(title="test-create-pr title", body="test-create-pr body")
     assert "Cannot make a PR from a detached HEAD" in str(exc_info.value)
+
+
+def test_codebase_create_pr_active_branch_is_default_branch(codebase: Codebase):
+    codebase.checkout(branch=codebase._op.default_branch)
+    codebase.files[0].remove()
+    codebase.commit()
+    with pytest.raises(ValueError) as exc_info:
+        codebase.create_pr(title="test-create-pr title", body="test-create-pr body")
+    assert "Cannot make a PR from the default branch" in str(exc_info.value)

tests/unit/codegen/extensions/lsp/test_workspace_sync.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 import pytest
 from lsprotocol.types import (
     DidChangeTextDocumentParams,
@@ -16,6 +18,7 @@
 from pytest_lsp import LanguageClient
 
 from codegen.sdk.core.codebase import Codebase
+from tests.unit.codegen.extensions.lsp.utils import apply_edit
 
 
 @pytest.fixture()
@@ -88,6 +91,7 @@ def example_function():
             """.strip(),
         ),
     ],
+    ids=["example_function"],
     indirect=["document_uri", "original"],
 )
 async def test_did_change(
@@ -137,7 +141,7 @@ def example_function():
     pass
                 """.strip(),
             },
-            "file://{worskpaceFolder}test.py",
+            "file://{worskpaceFolder}/test.py",
         ),
     ],
 )
@@ -165,11 +169,11 @@ async def test_did_close(
 
     # Verify the document is removed from the workspace
     document = await client.workspace_text_document_content_async(TextDocumentContentParams(uri=document_uri))
-    assert document.text == ""
+    assert document.text == original["test.py"]
 
 
 @pytest.mark.parametrize(
-    "original, document_uri, position, new_name, expected_text",
+    "original, document_uri, position, new_name, expected",
     [
         (
             {
@@ -182,15 +186,17 @@ def main():
                 """.strip(),
             },
             "file://{workspaceFolder}/test.py",
-            Position(line=0, character=0),  # Position of 'example_function'
+            Position(line=0, character=5),  # Position of 'example_function'
             "renamed_function",
-            """
+            {
+                "test.py": """
 def renamed_function():
     pass # modified
 
 def main():
     renamed_function()
                 """.strip(),
+            },
         ),
     ],
     indirect=["document_uri", "original"],
@@ -202,7 +208,7 @@ async def test_rename_after_sync(
     document_uri: str,
     position: Position,
     new_name: str,
-    expected_text: str,
+    assert_expected: Callable,
 ):
     # First open the document
     client.text_document_did_open(
@@ -243,8 +249,6 @@ async def test_rename_after_sync(
             new_name=new_name,
         )
     )
-
-    # Verify the rename was successful
-    document = await client.workspace_text_document_content_async(TextDocumentContentParams(uri=document_uri))
-    assert document is not None
-    assert document.text == expected_text
+    if result:
+        apply_edit(codebase, result)
+    assert_expected(codebase)

tests/unit/codegen/extensions/lsp/utils.py

@@ -10,7 +10,6 @@ def apply_edit(codebase: Codebase, edit: WorkspaceEdit):
             path = get_path(change.text_document.uri)
             file = codebase.get_file(str(path.relative_to(codebase.repo_path)))
             for edit in change.edits:
-                print("BRUH")
                 file.edit(edit.new_text)
         if isinstance(change, CreateFile):
             path = get_path(change.uri)