Skip to content

Commit 0f437b4

Browse files
authored
fix: make the codebase parse non-blocking; (#533)
# Motivation <!-- Why is this change necessary? --> This fixes an issue where the tool call for parsing the codebase, was blocked by codebase parse, this could cause the tool to timeout. Making it async let's the Agent check-in on when the parse is done # Content <!-- Please include a summary of the change --> # Testing <!-- How was the change tested? --> # Please check the following before marking your PR as ready for review - [ ] I have added tests for my changes - [ ] I have updated the documentation or added new documentation as needed
1 parent 0577888 commit 0f437b4

File tree

1 file changed

+5
-3
lines changed
  • codegen-examples/examples/codegen-mcp-server

1 file changed

+5
-3
lines changed

codegen-examples/examples/codegen-mcp-server/server.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,16 @@
99
class CodebaseState:
1010
"""Class to manage codebase state and parsing."""
1111

12-
parse_task: Optional[asyncio.Task] = None
12+
parse_task: Optional[asyncio.Future] = None
1313
parsed_codebase: Optional[Codebase] = None
14+
codebase_path: Optional[str] = None
1415
log_buffer: List[str] = field(default_factory=list)
1516

1617
async def parse(self, path: str) -> Codebase:
1718
"""Parse the codebase at the given path."""
1819
codebase = Codebase(path)
1920
self.parsed_codebase = codebase
21+
self.codebase_path = path
2022
return codebase
2123

2224
def reset(self) -> None:
@@ -46,7 +48,7 @@ def capture_output(*args, **kwargs) -> None:
4648
@mcp.tool(name="parse_codebase", description="Initiate codebase parsing")
4749
async def parse_codebase(codebase_path: Annotated[str, "path to the codebase to be parsed"]) -> Dict[str, str]:
4850
if not state.parse_task or state.parse_task.done():
49-
state.parse_task = asyncio.create_task(state.parse(codebase_path))
51+
state.parse_task = asyncio.get_event_loop().run_in_executor(None, lambda: state.parse(codebase_path))
5052
return {"message": "Codebase parsing initiated, this may take some time depending on the size of the codebase. Use the `check_parsing_status` tool to check if the parse has completed."}
5153
return {"message": "Codebase is already being parsed."}
5254

@@ -86,7 +88,7 @@ def main():
8688
print("starting codegen-mcp-server")
8789
run = mcp.run_stdio_async()
8890
print("codegen-mcp-server started")
89-
asyncio.run(run)
91+
asyncio.get_event_loop().run_until_complete(run)
9092

9193

9294
if __name__ == "__main__":

0 commit comments

Comments
 (0)