gabriel / muse public
agent_slots.py python
231 lines 7.4 KB
Raw
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ breaking 120 days ago
1 """Agent slot registry — track HD-derived agent accounts per hub.
2
3 Each hub maintains a monotonically-increasing account counter so that
4 new agent slots are assigned unique SLIP-0010 indices. Once assigned,
5 an agent's account index is permanent — it enables ERC8004 on-chain
6 binding and reproducible keypair derivation.
7
8 Storage
9 -------
10 ``~/.muse/agent-slots.toml``
11
12 TOML structure::
13
14 [localhost:1337]
15 next_account = 2
16
17 [localhost:1337.slots]
18 # name -> account index
19 orchestra = 1
20
21 [staging.musehub.ai]
22 next_account = 1
23
24 Slot 0 is reserved for the orchestrator identity on each hub.
25 Named agent slots start at account 1.
26 """
27
28 import contextlib
29 import fcntl
30 import os
31 import pathlib
32 import stat
33 import tempfile
34 from typing import Generator, TypedDict
35
36 try:
37 import tomllib
38 except ModuleNotFoundError: # Python < 3.11
39 import tomli as tomllib # type: ignore[no-reuse-def]
40
41 from muse.core.identity import hostname_from_url
42 from muse.core.paths import user_agent_slots_path as _user_agent_slots_path, user_muse_dir as _user_muse_dir
43
44 type _SlotNames = dict[str, int] # slot_name → account_index
45
46 class _HostBlock(TypedDict, total=False):
47 next_account: int
48 slots: _SlotNames
49
50 type _SlotsData = dict[str, _HostBlock]
51
52 _SLOTS_FILE = _user_agent_slots_path()
53 _SLOTS_DIR = _user_muse_dir()
54
55 class AgentSlot(TypedDict, total=False):
56 """A registered agent slot within an HD key hierarchy."""
57
58 name: str # human-readable slot name, e.g. "orchestra"
59 account: int # SLIP-0010 account index within DOMAIN_IDENTITY
60 hub: str # hostname this slot was registered against
61 msign_path: str # full SLIP-0010 path for the MSign key
62
63 # ---------------------------------------------------------------------------
64 # Internal helpers
65 # ---------------------------------------------------------------------------
66
67 def _toml_escape(s: str) -> str:
68 """Minimal TOML string escaping (backslash and double-quote)."""
69 return s.replace("\\", "\\\\").replace('"', '\\"')
70
71 def _load_raw(path: pathlib.Path) -> _SlotsData:
72 """Load the slots file as a raw dict. Returns empty dict if absent."""
73 if not path.is_file():
74 return {}
75 try:
76 with path.open("rb") as fh:
77 return tomllib.load(fh)
78 except Exception:
79 return {}
80
81 def _dump(data: _SlotsData) -> str:
82 """Serialise the slots dict back to TOML text."""
83 lines: list[str] = []
84 for hostname in sorted(data):
85 block = data[hostname]
86 if not isinstance(block, dict):
87 continue
88 lines.append(f'["{_toml_escape(hostname)}"]')
89 next_acct = block.get("next_account", 1)
90 lines.append(f"next_account = {int(next_acct)}")
91 slots: _SlotNames = block.get("slots", {})
92 if slots:
93 lines.append(f'["{_toml_escape(hostname)}".slots]')
94 for name in sorted(slots):
95 lines.append(f'{name} = {int(slots[name])}')
96 lines.append("")
97 return "\n".join(lines)
98
99 @contextlib.contextmanager
100 def _write_lock() -> Generator[None, None, None]:
101 """Exclusive advisory write-lock on the slots file."""
102 lock_path = _SLOTS_DIR / ".agent-slots.lock"
103 _SLOTS_DIR.mkdir(parents=True, exist_ok=True)
104 lock_fd = os.open(
105 str(lock_path),
106 os.O_CREAT | os.O_WRONLY | os.O_CLOEXEC,
107 stat.S_IRUSR | stat.S_IWUSR,
108 )
109 try:
110 fcntl.flock(lock_fd, fcntl.LOCK_EX)
111 yield
112 finally:
113 fcntl.flock(lock_fd, fcntl.LOCK_UN)
114 os.close(lock_fd)
115
116 def _save(data: _SlotsData, path: pathlib.Path) -> None:
117 """Atomically write the slots dict to *path* with mode 0o600."""
118 _SLOTS_DIR.mkdir(parents=True, exist_ok=True)
119 if path.is_symlink():
120 raise OSError(f"Security: {path} is a symlink — refusing to write.")
121 text = _dump(data)
122 fd, tmp_str = tempfile.mkstemp(dir=_SLOTS_DIR, prefix=".agent-slots-tmp-")
123 tmp = pathlib.Path(tmp_str)
124 try:
125 os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR)
126 with os.fdopen(fd, "w", encoding="utf-8") as fh:
127 fh.write(text)
128 os.replace(tmp, path)
129 except Exception:
130 tmp.unlink(missing_ok=True)
131 raise
132
133 # ---------------------------------------------------------------------------
134 # Public API
135 # ---------------------------------------------------------------------------
136
137 def get_next_account(hub_url: str) -> int:
138 """Return and atomically increment the next available agent account index.
139
140 Slot 0 is reserved for the orchestrator identity. Named agent
141 accounts start at 1 and are assigned sequentially — never reused.
142
143 Args:
144 hub_url: Hub URL or bare hostname.
145
146 Returns:
147 The next available account index (always >= 1).
148 """
149 hostname = hostname_from_url(hub_url)
150 with _write_lock():
151 data = _load_raw(_SLOTS_FILE)
152 block = data.setdefault(hostname, {})
153 current = int(block.get("next_account", 1))
154 if current < 1:
155 current = 1
156 block["next_account"] = current + 1
157 _save(data, _SLOTS_FILE)
158 return current
159
160 def register_slot(hub_url: str, name: str, account: int) -> AgentSlot:
161 """Register a named agent slot at the given account index.
162
163 If a slot with *name* already exists it is overwritten with the new
164 account index. Account indices must be unique across the hub — callers
165 should use :func:`get_next_account` to obtain a fresh index.
166
167 Args:
168 hub_url: Hub URL or bare hostname.
169 name: Human-readable slot label, e.g. ``"orchestra"``.
170 account: SLIP-0010 account index (>= 1; 0 is reserved for service).
171
172 Returns:
173 The registered :class:`AgentSlot`.
174 """
175 from muse.core.hdkeys import DOMAIN_IDENTITY, MUSE_PURPOSE # noqa: PLC0415
176 hostname = hostname_from_url(hub_url)
177 msign_path = f"m/{MUSE_PURPOSE}'/{DOMAIN_IDENTITY}'/1'/{account}'"
178 slot: AgentSlot = {
179 "name": name,
180 "account": account,
181 "hub": hostname,
182 "msign_path": msign_path,
183 }
184 with _write_lock():
185 data = _load_raw(_SLOTS_FILE)
186 block = data.setdefault(hostname, {})
187 slots = block.setdefault("slots", {})
188 slots[name] = account
189 _save(data, _SLOTS_FILE)
190 return slot
191
192 def list_slots(hub_url: str) -> list[AgentSlot]:
193 """Return all registered agent slots for *hub_url*, sorted by account index.
194
195 Args:
196 hub_url: Hub URL or bare hostname.
197
198 Returns:
199 List of :class:`AgentSlot` dicts, possibly empty.
200 """
201 from muse.core.hdkeys import DOMAIN_IDENTITY, MUSE_PURPOSE # noqa: PLC0415
202 hostname = hostname_from_url(hub_url)
203 data = _load_raw(_SLOTS_FILE)
204 block = data.get(hostname, {})
205 slots_raw: _SlotNames = block.get("slots", {})
206 result: list[AgentSlot] = []
207 for name, account in slots_raw.items():
208 msign_path = f"m/{MUSE_PURPOSE}'/{DOMAIN_IDENTITY}'/1'/{account}'"
209 result.append(AgentSlot(
210 name=name,
211 account=int(account),
212 hub=hostname,
213 msign_path=msign_path,
214 ))
215 result.sort(key=lambda s: s["account"])
216 return result
217
218 def peek_next_account(hub_url: str) -> int:
219 """Return the next account index without incrementing it.
220
221 Args:
222 hub_url: Hub URL or bare hostname.
223
224 Returns:
225 The next available account index (always >= 1).
226 """
227 hostname = hostname_from_url(hub_url)
228 data = _load_raw(_SLOTS_FILE)
229 block = data.get(hostname, {})
230 current = int(block.get("next_account", 1))
231 return max(current, 1)
File History 1 commit
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ 120 days ago