From cf99bdc51e94db98de2ff6b3c5923356fce9da97 Mon Sep 17 00:00:00 2001
From: Agent <agent@example.com>
Date: Thu, 16 Apr 2026 16:21:07 +0000
Subject: [PATCH] fix: add tomlkit to Dockerfile for comment-preserving TOML
 editing (#886)

---
 docker/agents/Dockerfile |  2 +-
 lib/hire-agent.sh        | 19 +++++++++----------
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/docker/agents/Dockerfile b/docker/agents/Dockerfile
index 2939230..1bcba89 100644
--- a/docker/agents/Dockerfile
+++ b/docker/agents/Dockerfile
@@ -2,7 +2,7 @@ FROM debian:bookworm-slim
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     bash curl git jq tmux python3 python3-pip openssh-client ca-certificates age shellcheck procps gosu \
-    && pip3 install --break-system-packages networkx \
+    && pip3 install --break-system-packages networkx tomlkit \
     && rm -rf /var/lib/apt/lists/*
 
 # Pre-built binaries (copied from docker/agents/bin/)
diff --git a/lib/hire-agent.sh b/lib/hire-agent.sh
index 45d0b0b..170389f 100644
--- a/lib/hire-agent.sh
+++ b/lib/hire-agent.sh
@@ -536,8 +536,7 @@ EOF
     echo "  Writing [agents.${section_name}] to ${toml_file}..."
     python3 -c '
 import sys
-import tomllib
-import tomli_w
+import tomlkit
 import re
 import pathlib
 
@@ -558,19 +557,19 @@ text = p.read_text()
 commented_pattern = rf"(?:^|\n)# \[agents\.{re.escape(section_name)}\](?:\n(?!# \[|\[)[^\n]*)*"
 text = re.sub(commented_pattern, "", text, flags=re.DOTALL)
 
-# Step 2: Parse TOML with tomllib
+# Step 2: Parse TOML with tomlkit (preserves comments and formatting)
 try:
-    data = tomllib.loads(text)
-except tomllib.TOMLDecodeError as e:
+    doc = tomlkit.parse(text)
+except Exception as e:
     print(f"Error: Invalid TOML in {toml_path}: {e}", file=sys.stderr)
     sys.exit(1)
 
 # Step 3: Ensure agents table exists
-if "agents" not in data:
-    data["agents"] = {}
+if "agents" not in doc:
+    doc.add("agents", tomlkit.table())
 
 # Step 4: Update the specific agent section
-data["agents"][section_name] = {
+doc["agents"][section_name] = {
     "base_url": base_url,
     "model": model,
     "api_key": "sk-no-key-required",
@@ -580,8 +579,8 @@ data["agents"][section_name] = {
     "poll_interval": int(poll_interval),
 }
 
-# Step 5: Serialize back to TOML
-output = tomli_w.dumps(data)
+# Step 5: Serialize back to TOML (preserves comments)
+output = tomlkit.dumps(doc)
 
 # Step 6: Write back
 p.write_text(output)