fix(subagents): harden roster lifecycle (#56)

Author: danielkov

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "kit"
-version = "0.1.109"
+version = "0.1.110"
 edition = "2024"
 rust-version = "1.94.0"
 publish = false
@@ -66,6 +66,7 @@ tower = { version = "=0.5.3", features = ["util"] }
 tracing = "=0.1.44"
 tracing-opentelemetry = { version = "=0.33.0", default-features = false }
 tracing-subscriber = { version = "=0.3.23", default-features = false, features = ["registry", "std"] }
+unicode-segmentation = "=1.13.3"
 unicode-width = "=0.2.2"
 url = "=2.5.8"
 uuid = { version = "=1.26.0", features = ["v5"] }

Cargo.toml

@@ -6,10 +6,29 @@
 import time
 
 write_lock = threading.Lock()
+log_lock = threading.Lock()
+state_lock = threading.Lock()
 next_session = 1
 supports_fork = "--no-fork" not in sys.argv
 supports_models = "--models" in sys.argv
+fail_first_close = "--fail-first-close" in sys.argv
+failed_close = False
 selected_models = {}
+
+
+def option(name):
+    prefix = name + "="
+    return next((arg[len(prefix):] for arg in sys.argv if arg.startswith(prefix)), None)
+
+
+request_log = option("--request-log")
+new_release = option("--new-release")
+fork_release = option("--fork-release")
+prompt_release = option("--prompt-release")
+prompt_release_text = option("--prompt-release-text")
+close_release = option("--close-release")
+close_release_session = option("--close-release-session")
+fail_close_session = option("--fail-close-session")
 model_ids = ["mock/default", "mock/requested"]
 
 if "--fail-start" in sys.argv:
@@ -26,11 +45,58 @@ def respond(request_id, result):
     send({"jsonrpc": "2.0", "id": request_id, "result": result})
 
 
+def log_request(request):
+    if request_log is None:
+        return
+    params = request.get("params", {})
+    entry = {"method": request.get("method")}
+    if "sessionId" in params:
+        entry["sessionId"] = params["sessionId"]
+    if request.get("method") == "session/prompt":
+        entry["text"] = params["prompt"][0]["text"]
+    with log_lock:
+        with open(request_log, "a", encoding="utf-8") as log:
+            log.write(json.dumps(entry, separators=(",", ":")) + "\n")
+            log.flush()
+
+
+def fork(request):
+    global next_session
+    if "--fail-fork" in sys.argv:
+        send({
+            "jsonrpc": "2.0",
+            "id": request["id"],
+            "error": {"code": -32000, "message": "fork failed"},
+        })
+        return
+    if not supports_fork:
+        send({
+            "jsonrpc": "2.0",
+            "id": request["id"],
+            "error": {"code": -32601, "message": "Method not found"},
+        })
+        return
+    source_id = request["params"]["sessionId"]
+    with state_lock:
+        session_id = f"branch-{next_session}"
+        next_session += 1
+        selected_models[session_id] = selected_models.get(source_id, model_ids[0])
+    while fork_release is not None and not os.path.exists(fork_release):
+        time.sleep(0.01)
+    respond(request["id"], {"sessionId": session_id})
+
+
 def prompt(request):
     params = request["params"]
     session_id = params["sessionId"]
     text = params["prompt"][0]["text"]
-    time.sleep(0.40)
+    should_gate = prompt_release is not None and (
+        prompt_release_text is None or prompt_release_text == text
+    )
+    while should_gate and not os.path.exists(prompt_release):
+        time.sleep(0.01)
+    if prompt_release is None:
+        time.sleep(0.40)
     if "MOCK_SELECTED_MODEL" in text:
         text = selected_models.get(session_id, model_ids[0])
     if "MOCK_STRUCTURED_OUTPUT" in text:
@@ -94,9 +160,36 @@ def prompt(request):
         os._exit(0)
 
 
+def close(request):
+    global failed_close
+    if "--slow-close" in sys.argv:
+        time.sleep(0.40)
+    session_id = request["params"]["sessionId"]
+    should_gate = close_release is not None and (
+        close_release_session is None or close_release_session == session_id
+    )
+    while should_gate and not os.path.exists(close_release):
+        time.sleep(0.01)
+    with state_lock:
+        should_fail = (fail_close_session == session_id and not failed_close) or (
+            fail_first_close and not failed_close
+        )
+        if should_fail:
+            failed_close = True
+    if should_fail:
+        send({
+            "jsonrpc": "2.0",
+            "id": request["id"],
+            "error": {"code": -32000, "message": "close failed"},
+        })
+    else:
+        respond(request["id"], {})
+
+
 for line in sys.stdin:
     request = json.loads(line)
     method = request.get("method")
+    log_request(request)
     if method == "initialize":
         respond(request["id"], {
             "protocolVersion": 1,
@@ -107,6 +200,8 @@ def prompt(request):
             },
         })
     elif method == "session/new":
+        while new_release is not None and not os.path.exists(new_release):
+            time.sleep(0.01)
         selected_models["base"] = model_ids[0]
         result = {"sessionId": "base"}
         if supports_models:
@@ -122,25 +217,7 @@ def prompt(request):
             }]
         respond(request["id"], result)
     elif method == "session/fork":
-        if "--fail-fork" in sys.argv:
-            send({
-                "jsonrpc": "2.0",
-                "id": request["id"],
-                "error": {"code": -32000, "message": "fork failed"},
-            })
-            continue
-        if not supports_fork:
-            send({
-                "jsonrpc": "2.0",
-                "id": request["id"],
-                "error": {"code": -32601, "message": "Method not found"},
-            })
-            continue
-        session_id = f"branch-{next_session}"
-        next_session += 1
-        source_id = request["params"]["sessionId"]
-        selected_models[session_id] = selected_models.get(source_id, model_ids[0])
-        respond(request["id"], {"sessionId": session_id})
+        threading.Thread(target=fork, args=(request,), daemon=True).start()
     elif method == "session/prompt":
         threading.Thread(target=prompt, args=(request,), daemon=True).start()
     elif method == "session/set_config_option":
@@ -158,6 +235,4 @@ def prompt(request):
     elif method == "session/cancel":
         pass
     elif method == "session/close":
-        if "--slow-close" in sys.argv:
-            time.sleep(0.40)
-        respond(request["id"], {})
+        threading.Thread(target=close, args=(request,), daemon=True).start()

fixtures/mock-acp.py

@@ -275,8 +275,7 @@ impl AcpHarnesses {
             command
                 .arg("--subagent-parent-id")
                 .arg(parent_id)
-                .arg("--subagent-parent-name")
-                .arg(parent_name);
+                .arg(format!("--subagent-parent-name={parent_name}"));
         }
         if resume {
             command.arg("--resume");
@@ -2000,10 +1999,18 @@ mod tests {
                 args.windows(2)
                     .any(|pair| pair == ["--subagent-parent-id", "s-parent"])
             );
-            assert!(
-                args.windows(2)
-                    .any(|pair| pair == ["--subagent-parent-name", "偵察 🦀"])
-            );
+            assert!(args.contains(&"--subagent-parent-name=偵察 🦀".into()));
+        }
+
+        #[test]
+        fn kit_child_parent_name_is_safe_when_it_starts_with_a_hyphen() {
+            let config = config(Some("s-parent"), Some("--reviewer"));
+            let command = config
+                .harnesses
+                .spawn(BUILTIN_HARNESS, &config, Some(("session", false)), 1)
+                .unwrap();
+
+            assert!(args(&command).contains(&"--subagent-parent-name=--reviewer".into()));
         }
 
         #[test]