Coverage for src/infra/hooks/locking.py: 53%

1"""Lock enforcement hooks for multi-agent file coordination.

2

3 Contains hooks for enforcing file locks when writing and for cleanup

4 on agent stop.

5"""

6

7 from __future__ import annotations

8

9 from collections.abc import Awaitable, Callable

10 from pathlib import Path

11 from typing import TYPE_CHECKING, Any

12

13 if TYPE_CHECKING:

14 from .dangerous_commands import PreToolUseHook

15

16 from ..tools.command_runner import run_command

17 from ..tools.env import SCRIPTS_DIR, get_lock_dir

18 from ..tools.locking import get_lock_holder

19 from .file_cache import FILE_PATH_KEYS, FILE_WRITE_TOOLS

20

21# Type alias for Stop hooks (using Any to avoid SDK import)

22 StopHook = Callable[

23 [Any, str | None, Any],

24 Awaitable[dict[str, Any]],

25]

26

27

28 def make_lock_enforcement_hook(

29 agent_id: str, repo_path: str | None = None

30) -> PreToolUseHook:

31 """Create a PreToolUse hook that enforces lock ownership for file writes.

32

33 Args:

34 agent_id: The agent ID to check lock ownership against.

35 repo_path: The repository root path, used as REPO_NAMESPACE for lock

36 key computation. Must match the REPO_NAMESPACE environment variable

37 set for the agent's shell scripts.

38

39 Returns:

40 An async hook function that blocks file writes unless the agent holds the lock.

41 """

42

43 async def enforce_lock_ownership(

44 hook_input: Any, # noqa: ANN401 - SDK type, avoid import

45 stderr: str | None,

46 context: Any, # noqa: ANN401 - SDK type, avoid import

47 ) -> dict[str, Any]:

48 """PreToolUse hook to block file writes unless this agent holds the lock."""

49 tool_name = hook_input["tool_name"]

50

51 # Only check file-write tools

52 if tool_name not in FILE_WRITE_TOOLS:

53 return {}

54

55 # Get the file path from the tool input

56 path_key = FILE_PATH_KEYS.get(tool_name)

57 if not path_key:

58 return {}

59

60 file_path = hook_input["tool_input"].get(path_key)

61 if not file_path:

62 # No path provided, can't check lock - allow (tool will fail anyway)

63 return {}

64

65 # Check if this agent holds the lock

66 # Pass repo_path as repo_namespace to match shell script key computation

67 lock_holder = get_lock_holder(file_path, repo_namespace=repo_path)

68

69 if lock_holder is None:

70 return {

71 "decision": "block",

72 "reason": f"File {file_path} is not locked. Acquire lock with: lock-try.sh {file_path}",

73 }

74

75 if lock_holder != agent_id:

76 return {

77 "decision": "block",

78 "reason": f"File {file_path} is locked by {lock_holder}. Wait or coordinate to acquire the lock.",

79 }

80

81 # Agent holds the lock, allow the write

82 return {}

83

84 return enforce_lock_ownership

85

86

87 def make_stop_hook(agent_id: str) -> StopHook:

88 """Create a Stop hook that cleans up locks for the given agent.

89

90 Args:

91 agent_id: The agent ID to clean up locks for when the agent stops.

92

93 Returns:

94 An async hook function suitable for use with ClaudeAgentOptions.hooks["Stop"].

95 """

96

97 async def cleanup_locks_on_stop(

98 hook_input: Any, # noqa: ANN401 - SDK type, avoid import

99 stderr: str | None,

100 context: Any, # noqa: ANN401 - SDK type, avoid import

101 ) -> dict[str, Any]:

102 """Stop hook to release all locks held by this agent."""

103 script = SCRIPTS_DIR / "lock-release-all.sh"

104 try:

105 run_command(

106 [str(script)],

107 cwd=Path.cwd(),

108 env={

109 "LOCK_DIR": str(get_lock_dir()),

110 "AGENT_ID": agent_id,

111 },

112 )

113 except Exception:

114 pass # Best effort cleanup, orchestrator has fallback

115 return {}

116

117 return cleanup_locks_on_stop