mirror of
https://github.com/shareAI-lab/analysis_claude_code.git
synced 2026-09-20 12:13:38 +08:00
757e84482f
The previous implementation always appended a new user message for background results regardless of the trailing message's role. After a tool round-trip (or right after the user query), the messages list could end up with two adjacent user messages, which is messy for prompt caching, token accounting and stricter API gateways. Now we merge the background-results block into the trailing user message when its role is already user, and only append a fresh user message otherwise. tool_result blocks still immediately follow their tool_use, so the protocol contract is preserved.
251 lines
9.8 KiB
Python
251 lines
9.8 KiB
Python
#!/usr/bin/env python3
|
|
# Harness: background execution -- the model thinks while the harness waits.
|
|
"""
|
|
s08_background_tasks.py - Background Tasks
|
|
|
|
Run commands in background threads. A notification queue is drained
|
|
before each LLM call to deliver results.
|
|
|
|
Main thread Background thread
|
|
+-----------------+ +-----------------+
|
|
| agent loop | | task executes |
|
|
| ... | | ... |
|
|
| [LLM call] <---+------- | enqueue(result) |
|
|
| ^drain queue | +-----------------+
|
|
+-----------------+
|
|
|
|
Timeline:
|
|
Agent ----[spawn A]----[spawn B]----[other work]----
|
|
| |
|
|
v v
|
|
[A runs] [B runs] (parallel)
|
|
| |
|
|
+-- notification queue --> [results injected]
|
|
|
|
Key insight: "Fire and forget -- the agent doesn't block while the command runs."
|
|
"""
|
|
|
|
import os
|
|
import subprocess
|
|
import threading
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from anthropic import Anthropic
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv(override=True)
|
|
|
|
if os.getenv("ANTHROPIC_BASE_URL"):
|
|
os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)
|
|
|
|
WORKDIR = Path.cwd()
|
|
client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
|
|
MODEL = os.environ["MODEL_ID"]
|
|
|
|
SYSTEM = f"You are a coding agent at {WORKDIR}. Use background_run for long-running commands."
|
|
|
|
|
|
# -- BackgroundManager: threaded execution + notification queue --
|
|
class BackgroundManager:
|
|
def __init__(self):
|
|
self.tasks = {} # task_id -> {status, result, command}
|
|
self._notification_queue = [] # completed task results
|
|
self._lock = threading.Lock()
|
|
|
|
def run(self, command: str) -> str:
|
|
"""Start a background thread, return task_id immediately."""
|
|
task_id = str(uuid.uuid4())[:8]
|
|
self.tasks[task_id] = {"status": "running", "result": None, "command": command}
|
|
thread = threading.Thread(
|
|
target=self._execute, args=(task_id, command), daemon=True
|
|
)
|
|
thread.start()
|
|
return f"Background task {task_id} started: {command[:80]}"
|
|
|
|
def _execute(self, task_id: str, command: str):
|
|
"""Thread target: run subprocess, capture output, push to queue."""
|
|
try:
|
|
r = subprocess.run(
|
|
command, shell=True, cwd=WORKDIR,
|
|
capture_output=True, text=True, timeout=300
|
|
)
|
|
output = (r.stdout + r.stderr).strip()[:50000]
|
|
status = "completed"
|
|
except subprocess.TimeoutExpired:
|
|
output = "Error: Timeout (300s)"
|
|
status = "timeout"
|
|
except Exception as e:
|
|
output = f"Error: {e}"
|
|
status = "error"
|
|
self.tasks[task_id]["status"] = status
|
|
self.tasks[task_id]["result"] = output or "(no output)"
|
|
with self._lock:
|
|
self._notification_queue.append({
|
|
"task_id": task_id,
|
|
"status": status,
|
|
"command": command[:80],
|
|
"result": (output or "(no output)")[:500],
|
|
})
|
|
|
|
def check(self, task_id: str = None) -> str:
|
|
"""Check status of one task or list all."""
|
|
if task_id:
|
|
t = self.tasks.get(task_id)
|
|
if not t:
|
|
return f"Error: Unknown task {task_id}"
|
|
return f"[{t['status']}] {t['command'][:60]}\n{t.get('result') or '(running)'}"
|
|
lines = []
|
|
for tid, t in self.tasks.items():
|
|
lines.append(f"{tid}: [{t['status']}] {t['command'][:60]}")
|
|
return "\n".join(lines) if lines else "No background tasks."
|
|
|
|
def drain_notifications(self) -> list:
|
|
"""Return and clear all pending completion notifications."""
|
|
with self._lock:
|
|
notifs = list(self._notification_queue)
|
|
self._notification_queue.clear()
|
|
return notifs
|
|
|
|
|
|
BG = BackgroundManager()
|
|
|
|
|
|
# -- Tool implementations --
|
|
def safe_path(p: str) -> Path:
|
|
path = (WORKDIR / p).resolve()
|
|
if not path.is_relative_to(WORKDIR):
|
|
raise ValueError(f"Path escapes workspace: {p}")
|
|
return path
|
|
|
|
def run_bash(command: str) -> str:
|
|
dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"]
|
|
if any(d in command for d in dangerous):
|
|
return "Error: Dangerous command blocked"
|
|
try:
|
|
r = subprocess.run(command, shell=True, cwd=WORKDIR,
|
|
capture_output=True, text=True, timeout=120)
|
|
out = (r.stdout + r.stderr).strip()
|
|
return out[:50000] if out else "(no output)"
|
|
except subprocess.TimeoutExpired:
|
|
return "Error: Timeout (120s)"
|
|
|
|
def run_read(path: str, limit: int = None) -> str:
|
|
try:
|
|
lines = safe_path(path).read_text().splitlines()
|
|
if limit and limit < len(lines):
|
|
lines = lines[:limit] + [f"... ({len(lines) - limit} more)"]
|
|
return "\n".join(lines)[:50000]
|
|
except Exception as e:
|
|
return f"Error: {e}"
|
|
|
|
def run_write(path: str, content: str) -> str:
|
|
try:
|
|
fp = safe_path(path)
|
|
fp.parent.mkdir(parents=True, exist_ok=True)
|
|
fp.write_text(content)
|
|
return f"Wrote {len(content)} bytes"
|
|
except Exception as e:
|
|
return f"Error: {e}"
|
|
|
|
def run_edit(path: str, old_text: str, new_text: str) -> str:
|
|
try:
|
|
fp = safe_path(path)
|
|
c = fp.read_text()
|
|
if old_text not in c:
|
|
return f"Error: Text not found in {path}"
|
|
fp.write_text(c.replace(old_text, new_text, 1))
|
|
return f"Edited {path}"
|
|
except Exception as e:
|
|
return f"Error: {e}"
|
|
|
|
|
|
TOOL_HANDLERS = {
|
|
"bash": lambda **kw: run_bash(kw["command"]),
|
|
"read_file": lambda **kw: run_read(kw["path"], kw.get("limit")),
|
|
"write_file": lambda **kw: run_write(kw["path"], kw["content"]),
|
|
"edit_file": lambda **kw: run_edit(kw["path"], kw["old_text"], kw["new_text"]),
|
|
"background_run": lambda **kw: BG.run(kw["command"]),
|
|
"check_background": lambda **kw: BG.check(kw.get("task_id")),
|
|
}
|
|
|
|
TOOLS = [
|
|
{"name": "bash", "description": "Run a shell command (blocking).",
|
|
"input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}},
|
|
{"name": "read_file", "description": "Read file contents.",
|
|
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "limit": {"type": "integer"}}, "required": ["path"]}},
|
|
{"name": "write_file", "description": "Write content to file.",
|
|
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}},
|
|
{"name": "edit_file", "description": "Replace exact text in file.",
|
|
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}}, "required": ["path", "old_text", "new_text"]}},
|
|
{"name": "background_run", "description": "Run command in background thread. Returns task_id immediately.",
|
|
"input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}},
|
|
{"name": "check_background", "description": "Check background task status. Omit task_id to list all.",
|
|
"input_schema": {"type": "object", "properties": {"task_id": {"type": "string"}}}},
|
|
]
|
|
|
|
|
|
def agent_loop(messages: list):
|
|
while True:
|
|
# Drain background notifications and inject before the next LLM call.
|
|
# Merge into the trailing user message when possible to avoid emitting
|
|
# two consecutive user messages (which is messy for caching/debugging).
|
|
notifs = BG.drain_notifications()
|
|
if notifs and messages:
|
|
notif_text = "\n".join(
|
|
f"[bg:{n['task_id']}] {n['status']}: {n['result']}" for n in notifs
|
|
)
|
|
bg_block = {
|
|
"type": "text",
|
|
"text": f"<background-results>\n{notif_text}\n</background-results>",
|
|
}
|
|
last = messages[-1]
|
|
if last["role"] == "user":
|
|
if isinstance(last["content"], str):
|
|
last["content"] = [
|
|
{"type": "text", "text": last["content"]},
|
|
bg_block,
|
|
]
|
|
else:
|
|
last["content"] = list(last["content"]) + [bg_block]
|
|
else:
|
|
messages.append({"role": "user", "content": [bg_block]})
|
|
response = client.messages.create(
|
|
model=MODEL, system=SYSTEM, messages=messages,
|
|
tools=TOOLS, max_tokens=8000,
|
|
)
|
|
messages.append({"role": "assistant", "content": response.content})
|
|
if response.stop_reason != "tool_use":
|
|
return
|
|
results = []
|
|
for block in response.content:
|
|
if block.type == "tool_use":
|
|
handler = TOOL_HANDLERS.get(block.name)
|
|
try:
|
|
output = handler(**block.input) if handler else f"Unknown tool: {block.name}"
|
|
except Exception as e:
|
|
output = f"Error: {e}"
|
|
print(f"> {block.name}:")
|
|
print(str(output)[:200])
|
|
results.append({"type": "tool_result", "tool_use_id": block.id, "content": str(output)})
|
|
messages.append({"role": "user", "content": results})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
history = []
|
|
while True:
|
|
try:
|
|
query = input("\033[36ms08 >> \033[0m")
|
|
except (EOFError, KeyboardInterrupt):
|
|
break
|
|
if query.strip().lower() in ("q", "exit", ""):
|
|
break
|
|
history.append({"role": "user", "content": query})
|
|
agent_loop(history)
|
|
response_content = history[-1]["content"]
|
|
if isinstance(response_content, list):
|
|
for block in response_content:
|
|
if hasattr(block, "text"):
|
|
print(block.text)
|
|
print()
|