@@ -154,19 +154,33 @@ def register_checkpoint(
154154 snapshot_path : str ,
155155 target_path : str ,
156156 tool : str = "" ,
157+ session_id : str | None = None ,
158+ user_id : str | None = None ,
159+ workspace_id : str | None = None ,
157160) -> str :
158- """Register a snapshot and return opaque id (thread + cross-process locked)."""
161+ """Register a snapshot and return opaque id (thread + cross-process locked).
162+
163+ Optional ownership fields bind restore to the creating session/user.
164+ """
159165 root = _project_root ()
160166 cid = uuid .uuid4 ().hex
161167
162168 def _do () -> None :
163169 reg = _load_registry (root )
164- reg [ cid ] = {
170+ entry : dict [ str , Any ] = {
165171 "snapshot" : str (snapshot_path ),
166172 "target" : str (target_path ),
167173 "tool" : tool ,
168174 "created_at" : datetime .now (timezone .utc ).isoformat (),
175+ "workspace_root" : str (root ),
169176 }
177+ if session_id :
178+ entry ["session_id" ] = str (session_id )
179+ if user_id :
180+ entry ["user_id" ] = str (user_id )
181+ if workspace_id :
182+ entry ["workspace_id" ] = str (workspace_id )
183+ reg [cid ] = entry
170184 if len (reg ) > 400 :
171185 items = sorted (
172186 reg .items (),
@@ -248,10 +262,22 @@ def snapshot_path_for_tool(name: str, arguments: dict[str, Any]) -> str | None:
248262 idx = dest_root / "INDEX.txt"
249263 with idx .open ("a" , encoding = "utf-8" ) as f :
250264 f .write (f"{ name } \t { target } \t { dest } \n " )
265+ _sid = str (
266+ (arguments or {}).get ("_session_id" )
267+ or (arguments or {}).get ("session_id" )
268+ or ""
269+ ).strip () or None
270+ _uid = str (
271+ (arguments or {}).get ("_user_id" )
272+ or (arguments or {}).get ("user_id" )
273+ or ""
274+ ).strip () or None
251275 cid = register_checkpoint (
252276 snapshot_path = str (dest ),
253277 target_path = str (target ),
254278 tool = name ,
279+ session_id = _sid ,
280+ user_id = _uid ,
255281 )
256282 # also write id into index for operators
257283 with idx .open ("a" , encoding = "utf-8" ) as f :
@@ -267,12 +293,18 @@ def list_recent_checkpoints(limit: int = 20) -> list[str]:
267293 return [str (p ) for p in dirs [:limit ]]
268294
269295
270- def restore_checkpoint_file (snapshot_path_or_id : str ) -> dict [str , Any ]:
296+ def restore_checkpoint_file (
297+ snapshot_path_or_id : str ,
298+ * ,
299+ session_id : str | None = None ,
300+ user_id : str | None = None ,
301+ ) -> dict [str , Any ]:
271302 """Restore by opaque id (preferred) or legacy snapshot path.
272303
273304 Security:
274305 - registry targets must stay under project root
275306 - legacy path snapshots must live under .tevarn/checkpoints/
307+ - if registry has session_id/user_id, caller must match when provided
276308 """
277309 raw = str (snapshot_path_or_id or "" ).strip ()
278310 if not raw :
@@ -284,6 +316,13 @@ def restore_checkpoint_file(snapshot_path_or_id: str) -> dict[str, Any]:
284316 # Opaque id path
285317 entry = lookup_checkpoint (raw )
286318 if entry and entry .get ("backend" ) == "python" and entry .get ("snapshot" ):
319+ # Ownership gate (shared-deploy safety)
320+ reg_sid = str (entry .get ("session_id" ) or "" ).strip ()
321+ reg_uid = str (entry .get ("user_id" ) or "" ).strip ()
322+ if reg_sid and session_id and str (session_id ).strip () != reg_sid :
323+ return {"ok" : False , "error" : "checkpoint session mismatch" }
324+ if reg_uid and user_id and str (user_id ).strip () != reg_uid :
325+ return {"ok" : False , "error" : "checkpoint owner mismatch" }
287326 snap = Path (str (entry ["snapshot" ])).expanduser ()
288327 target = Path (str (entry .get ("target" ) or "" )).expanduser ()
289328 try :
0 commit comments