Merge pull request #1129 from Codium-ai/tr/err_protections

Tr/err protections
Codium-ai · Aug 13, 2024 · 1f7a8ea · 1f7a8ea
2 parents d8572f8 + 38638bd
commit 1f7a8ea
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 4 deletions.
diff --git a/pr_agent/git_providers/bitbucket_provider.py b/pr_agent/git_providers/bitbucket_provider.py
@@ -46,6 +46,7 @@ def __init__(
         self.temp_comments = []
         self.incremental = incremental
         self.diff_files = None
+        self.git_files = None
         if pr_url:
             self.set_pr(pr_url)
         self.bitbucket_comment_api_url = self.pr._BitbucketBase__data["links"]["comments"]["href"]

diff --git a/pr_agent/secret_providers/google_cloud_storage_secret_provider.py b/pr_agent/secret_providers/google_cloud_storage_secret_provider.py
@@ -22,7 +22,7 @@ def get_secret(self, secret_name: str) -> str:
             blob = self.bucket.blob(secret_name)
             return blob.download_as_string()
         except Exception as e:
-            get_logger().error(f"Failed to get secret {secret_name} from Google Cloud Storage: {e}")
+            get_logger().warning(f"Failed to get secret {secret_name} from Google Cloud Storage: {e}")
             return ""
 
     def store_secret(self, secret_name: str, secret_value: str):

diff --git a/pr_agent/servers/gitlab_webhook.py b/pr_agent/servers/gitlab_webhook.py
@@ -87,6 +87,10 @@ async def inner(data: dict):
         if request.headers.get("X-Gitlab-Token") and secret_provider:
             request_token = request.headers.get("X-Gitlab-Token")
             secret = secret_provider.get_secret(request_token)
+            if not secret:
+                get_logger().warning(f"Empty secret retrieved, request_token: {request_token}")
+                return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED,
+                                    content=jsonable_encoder({"message": "unauthorized"}))
             try:
                 secret_dict = json.loads(secret)
                 gitlab_token = secret_dict["gitlab_token"]

diff --git a/pr_agent/tools/pr_code_suggestions.py b/pr_agent/tools/pr_code_suggestions.py
@@ -450,8 +450,24 @@ def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet):
             original_initial_line = None
             for file in self.diff_files:
                 if file.filename.strip() == relevant_file:
-                    if file.head_file:  # in bitbucket, head_file is empty. toDo: fix this
-                        original_initial_line = file.head_file.splitlines()[relevant_lines_start - 1]
+                    if file.head_file:
+                        file_lines = file.head_file.splitlines()
+                        if relevant_lines_start > len(file_lines):
+                            get_logger().warning(
+                                "Could not dedent code snippet, because relevant_lines_start is out of range",
+                                artifact={'filename': file.filename,
+                                          'file_content': file.head_file,
+                                          'relevant_lines_start': relevant_lines_start,
+                                          'new_code_snippet': new_code_snippet})
+                            return new_code_snippet
+                        else:
+                            original_initial_line = file_lines[relevant_lines_start - 1]
+                    else:
+                        get_logger().warning("Could not dedent code snippet, because head_file is missing",
+                                             artifact={'filename': file.filename,
+                                                       'relevant_lines_start': relevant_lines_start,
+                                                       'new_code_snippet': new_code_snippet})
+                        return new_code_snippet
                     break
             if original_initial_line:
                 suggested_initial_line = new_code_snippet.splitlines()[0]
@@ -461,7 +477,7 @@ def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet):
                 if delta_spaces > 0:
                     new_code_snippet = textwrap.indent(new_code_snippet, delta_spaces * " ").rstrip('\n')
         except Exception as e:
-            get_logger().error(f"Could not dedent code snippet for file {relevant_file}, error: {e}")
+            get_logger().error(f"Error when dedenting code snippet for file {relevant_file}, error: {e}")
 
         return new_code_snippet
 

diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py
@@ -510,6 +510,12 @@ def _prepare_file_labels(self):
         file_label_dict = {}
         for file in self.data['pr_files']:
             try:
+                required_fields = ['changes_summary', 'changes_title', 'filename', 'label']
+                if not all(field in file for field in required_fields):
+                    # can happen for example if a YAML generation was interrupted in the middle (no more tokens)
+                    get_logger().warning(f"Missing required fields in file label dict {self.pr_id}, skipping file",
+                                         artifact={"file": file})
+                    continue
                 filename = file['filename'].replace("'", "`").replace('"', '`')
                 changes_summary = file['changes_summary']
                 changes_title = file['changes_title'].strip()