Makes apply_sync_write, write_to_local_delta, and store modify truly async, preventing the mcp-memory-server from locking the tokio executor during disk IO
45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
import re
|
|
|
|
def fix_handlers(filepath):
|
|
with open(filepath, "r") as f:
|
|
content = f.read()
|
|
|
|
# Find the positions of self.state.write_to_local_delta( and self.state.apply_sync_write(
|
|
patterns = ["self.state.write_to_local_delta(|", "self.state.apply_sync_write(|"]
|
|
|
|
# Actually, we can just find 'write_to_local_delta(|' and 'apply_sync_write(|'
|
|
|
|
for method in ["write_to_local_delta", "apply_sync_write"]:
|
|
start_idx = 0
|
|
while True:
|
|
idx = content.find(f"self.state.{method}(|", start_idx)
|
|
if idx == -1:
|
|
break
|
|
|
|
# Find the closing bracket that matches the first `{` after `|`
|
|
brace_idx = content.find("{", idx)
|
|
|
|
# count braces
|
|
brace_count = 1
|
|
curr_idx = brace_idx + 1
|
|
while brace_count > 0 and curr_idx < len(content):
|
|
if content[curr_idx] == '{':
|
|
brace_count += 1
|
|
elif content[curr_idx] == '}':
|
|
brace_count -= 1
|
|
curr_idx += 1
|
|
|
|
# curr_idx is now right after the matching '}'
|
|
# the original code was '});'
|
|
if content[curr_idx:curr_idx+2] == ");":
|
|
# we want to insert .await
|
|
content = content[:curr_idx+1] + ".await" + content[curr_idx+1:]
|
|
start_idx = curr_idx + 8 # move past `.await);`
|
|
else:
|
|
start_idx = curr_idx
|
|
|
|
with open(filepath, "w") as f:
|
|
f.write(content)
|
|
|
|
fix_handlers("C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\server\\src\\handlers.rs")
|