Skip to content

Commit 93cdfb2

Browse files
committed
fix: add asyncio lock to context file writes
On Windows, concurrent async writes to the context.jsonl file can cause PermissionError due to file locking semantics. Multiple coroutines (append_message, update_token_count, checkpoint) may attempt to write simultaneously. Add an asyncio.Lock to serialize all file write operations in the Context class, preventing concurrent access to the context file. Fixes #1429
1 parent 7ba9695 commit 93cdfb2

1 file changed

Lines changed: 12 additions & 7 deletions

File tree

src/kimi_cli/soul/context.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import asyncio
34
import json
45
from collections.abc import Sequence
56
from pathlib import Path
@@ -21,6 +22,7 @@ def __init__(self, file_backend: Path):
2122
self._next_checkpoint_id: int = 0
2223
"""The ID of the next checkpoint, starting from 0, incremented after each checkpoint."""
2324
self._system_prompt: str | None = None
25+
self._write_lock = asyncio.Lock()
2426

2527
async def restore(self) -> bool:
2628
logger.debug("Restoring context from file: {file_backend}", file_backend=self._file_backend)
@@ -105,8 +107,9 @@ async def checkpoint(self, add_user_message: bool):
105107
self._next_checkpoint_id += 1
106108
logger.debug("Checkpointing, ID: {id}", id=checkpoint_id)
107109

108-
async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
109-
await f.write(json.dumps({"role": "_checkpoint", "id": checkpoint_id}) + "\n")
110+
async with self._write_lock:
111+
async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
112+
await f.write(json.dumps({"role": "_checkpoint", "id": checkpoint_id}) + "\n")
110113
if add_user_message:
111114
await self.append_message(
112115
Message(role="user", content=[system(f"CHECKPOINT {checkpoint_id}")])
@@ -203,13 +206,15 @@ async def append_message(self, message: Message | Sequence[Message]):
203206
messages = [message] if isinstance(message, Message) else message
204207
self._history.extend(messages)
205208

206-
async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
207-
for message in messages:
208-
await f.write(message.model_dump_json(exclude_none=True) + "\n")
209+
async with self._write_lock:
210+
async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
211+
for message in messages:
212+
await f.write(message.model_dump_json(exclude_none=True) + "\n")
209213

210214
async def update_token_count(self, token_count: int):
211215
logger.debug("Updating token count in context: {token_count}", token_count=token_count)
212216
self._token_count = token_count
213217

214-
async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
215-
await f.write(json.dumps({"role": "_usage", "token_count": token_count}) + "\n")
218+
async with self._write_lock:
219+
async with aiofiles.open(self._file_backend, "a", encoding="utf-8") as f:
220+
await f.write(json.dumps({"role": "_usage", "token_count": token_count}) + "\n")

0 commit comments

Comments
 (0)