fix: guard against missing key in annotation dependency tracking

The `_extract_names_from_annotation` method accessed
`self.definitions[self.current_top_level_name]` without verifying
that `current_top_level_name` exists in the definitions dict. When
visiting methods inside built-in types (e.g., `str.removeprefix`),
the dotted name was set as `current_top_level_name` but never added
to `definitions`, causing a KeyError on CI.

Author: KRRT7

codeflash/languages/python/context/unused_definition_remover.py

@@ -202,7 +202,12 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
     def _extract_names_from_annotation(self, node: cst.CSTNode) -> None:
         if isinstance(node, cst.Name):
             name = node.value
-            if name in self.definitions and name != self.current_top_level_name and self.current_top_level_name:
+            if (
+                name in self.definitions
+                and name != self.current_top_level_name
+                and self.current_top_level_name
+                and self.current_top_level_name in self.definitions
+            ):
                 self.definitions[self.current_top_level_name].dependencies.add(name)
         elif isinstance(node, cst.Subscript):
             self._extract_names_from_annotation(node.value)

pyproject.toml

@@ -225,7 +225,7 @@ module = ["jedi", "jedi.api.classes", "inquirer", "inquirer.themes", "numba", "d
 ignore_missing_imports = true
 
 [[tool.mypy.overrides]]
-module = ["codeflash.code_utils.codeflash_wrap_decorator", "codeflash.code_utils.instrument_existing_tests", "codeflash.optimization.optimizer"]
+module = ["codeflash.code_utils.codeflash_wrap_decorator", "codeflash.code_utils.instrument_existing_tests", "codeflash.optimization.optimizer", "codeflash.languages.python.context.unused_definition_remover"]
 disable_error_code = ["attr-defined", "return-value", "no-untyped-call", "no-untyped-def", "arg-type", "assignment", "var-annotated", "no-any-return", "call-overload", "union-attr", "unreachable", "list-item"]
 
 [tool.pydantic-mypy]