Cyber Catalyst Team commited on
Commit
4c548f9
·
1 Parent(s): d946914

Implement non-blocking async tool execution and context auto-compaction for 24/7 continuous operation

Browse files
Files changed (1) hide show
  1. backend.py +103 -28
backend.py CHANGED
@@ -247,8 +247,8 @@ def repair_arguments(func_name: str, args: dict) -> tuple[dict, list[str]]:
247
  return repaired_args, notes
248
 
249
 
250
- def execute_tool(name: str, arguments: dict) -> str:
251
- """Execute a tool and return its output as a string."""
252
  try:
253
  if name == "read_file":
254
  path = _safe_path(arguments["path"])
@@ -273,24 +273,35 @@ def execute_tool(name: str, arguments: dict) -> str:
273
  blocked = ["rm -rf /", "mkfs", "dd if=", ":(){", "fork bomb"]
274
  if any(b in command.lower() for b in blocked):
275
  return "Error: Command blocked for safety reasons"
276
- result = subprocess.run(
277
- ["bash", "-c", command],
 
 
 
 
278
  cwd=WORKSPACE_DIR,
279
- capture_output=True,
280
- text=True,
281
- timeout=30,
282
  env={**os.environ, "HOME": "/tmp", "PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")},
283
  )
284
- output = ""
285
- if result.stdout:
286
- output += result.stdout
287
- if result.stderr:
288
- output += ("\n" if output else "") + f"[stderr] {result.stderr}"
289
- if result.returncode != 0:
290
- output += f"\n[exit code: {result.returncode}]"
291
- if not output:
292
- output = "[command completed with no output]"
293
- # Truncate very long outputs
 
 
 
 
 
 
 
 
 
 
294
  if len(output) > 20000:
295
  output = output[:20000] + f"\n\n[Truncated — output is {len(output)} chars]"
296
  return output
@@ -320,14 +331,26 @@ def execute_tool(name: str, arguments: dict) -> str:
320
  pattern = arguments["pattern"]
321
  search_path = arguments.get("path", ".")
322
  path = _safe_path(search_path)
323
- result = subprocess.run(
324
- ["grep", "-rn", "--include=*", pattern, str(path)],
325
- capture_output=True,
326
- text=True,
327
- timeout=10,
 
 
328
  cwd=WORKSPACE_DIR,
329
  )
330
- output = result.stdout if result.stdout else "No matches found"
 
 
 
 
 
 
 
 
 
 
331
  if len(output) > 10000:
332
  output = output[:10000] + "\n\n[Truncated]"
333
  return output
@@ -335,10 +358,6 @@ def execute_tool(name: str, arguments: dict) -> str:
335
  else:
336
  return f"Error: Unknown tool: {name}"
337
 
338
- except subprocess.TimeoutExpired:
339
- return "Error: Command timed out after 30 seconds"
340
- except ValueError as e:
341
- return f"Error: {str(e)}"
342
  except Exception as e:
343
  return f"Error executing {name}: {str(e)}"
344
 
@@ -572,6 +591,57 @@ IMPORTANT RULES:
572
  """
573
 
574
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
575
  @app.post("/v1/chat/completions")
576
  async def chat_completions(request: Request, authorization: str = Header(None)):
577
  auth(authorization)
@@ -611,6 +681,9 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
611
  else:
612
  final_messages = list(messages)
613
 
 
 
 
614
  # Save the user's message to DB
615
  user_msg = next((m for m in reversed(messages) if m.get("role") == "user"), None)
616
  if user_msg:
@@ -646,6 +719,8 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
646
  async with completions_semaphore:
647
  try:
648
  for round_num in range(MAX_TOOL_ROUNDS + 1):
 
 
649
  kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
650
  if is_agentic:
651
  kwargs["tools"] = TOOLS
@@ -757,7 +832,7 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
757
  yield make_chunk(request_id, requested_model, "\n")
758
 
759
  # Execute the tool
760
- result = execute_tool(func_name, repaired_args)
761
 
762
  # Append teaching note if repaired
763
  if repair_notes:
 
247
  return repaired_args, notes
248
 
249
 
250
+ async def execute_tool(name: str, arguments: dict) -> str:
251
+ """Execute a tool and return its output as a string asynchronously."""
252
  try:
253
  if name == "read_file":
254
  path = _safe_path(arguments["path"])
 
273
  blocked = ["rm -rf /", "mkfs", "dd if=", ":(){", "fork bomb"]
274
  if any(b in command.lower() for b in blocked):
275
  return "Error: Command blocked for safety reasons"
276
+
277
+ # ASYNC SUBPROCESS - This prevents the FastAPI server from freezing!
278
+ process = await asyncio.create_subprocess_shell(
279
+ command,
280
+ stdout=asyncio.subprocess.PIPE,
281
+ stderr=asyncio.subprocess.PIPE,
282
  cwd=WORKSPACE_DIR,
 
 
 
283
  env={**os.environ, "HOME": "/tmp", "PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")},
284
  )
285
+
286
+ try:
287
+ stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=30)
288
+ output = ""
289
+ if stdout:
290
+ output += stdout.decode('utf-8', errors='replace')
291
+ if stderr:
292
+ output += ("\n" if output else "") + f"[stderr] {stderr.decode('utf-8', errors='replace')}"
293
+ if process.returncode != 0:
294
+ output += f"\n[exit code: {process.returncode}]"
295
+ if not output:
296
+ output = "[command completed with no output]"
297
+ except asyncio.TimeoutError:
298
+ try:
299
+ process.kill()
300
+ except Exception:
301
+ pass
302
+ await process.communicate()
303
+ return "Error: Command timed out after 30 seconds"
304
+
305
  if len(output) > 20000:
306
  output = output[:20000] + f"\n\n[Truncated — output is {len(output)} chars]"
307
  return output
 
331
  pattern = arguments["pattern"]
332
  search_path = arguments.get("path", ".")
333
  path = _safe_path(search_path)
334
+
335
+ # Escape single quotes in pattern for safety
336
+ escaped_pattern = pattern.replace("'", "'\\''")
337
+ process = await asyncio.create_subprocess_shell(
338
+ f"grep -rn --include=* '{escaped_pattern}' '{path}'",
339
+ stdout=asyncio.subprocess.PIPE,
340
+ stderr=asyncio.subprocess.PIPE,
341
  cwd=WORKSPACE_DIR,
342
  )
343
+ try:
344
+ stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=10)
345
+ output = stdout.decode('utf-8', errors='replace') if stdout else "No matches found"
346
+ except asyncio.TimeoutError:
347
+ try:
348
+ process.kill()
349
+ except Exception:
350
+ pass
351
+ await process.communicate()
352
+ return "Error: Grep search timed out"
353
+
354
  if len(output) > 10000:
355
  output = output[:10000] + "\n\n[Truncated]"
356
  return output
 
358
  else:
359
  return f"Error: Unknown tool: {name}"
360
 
 
 
 
 
361
  except Exception as e:
362
  return f"Error executing {name}: {str(e)}"
363
 
 
591
  """
592
 
593
 
594
+ def compact_history(messages: list) -> list:
595
+ """
596
+ Compact conversation history to prevent context window overflow.
597
+ Replaces massive tool call outputs with concise summaries if history is long.
598
+ """
599
+ # Only compact if messages count exceeds 15 (to maintain normal conversation)
600
+ if len(messages) <= 15:
601
+ return messages
602
+
603
+ compacted = []
604
+ # Always keep the system prompt (typically the first message)
605
+ if messages and messages[0].get("role") == "system":
606
+ compacted.append(messages[0])
607
+ start_idx = 1
608
+ else:
609
+ start_idx = 0
610
+
611
+ # Keep the last 4 messages exactly as they are to preserve immediate context
612
+ recent_count = 4
613
+ mid_messages = messages[start_idx:-recent_count]
614
+ recent_messages = messages[-recent_count:]
615
+
616
+ for msg in mid_messages:
617
+ role = msg.get("role")
618
+ content = msg.get("content") or ""
619
+
620
+ if role == "tool":
621
+ # Compress massive tool outputs (like bash stdout or file reads)
622
+ if len(content) > 1000:
623
+ summary = f"[Tool output compacted: {content[:200]}... (Total {len(content)} chars truncated for context preservation)]"
624
+ compacted.append({
625
+ "role": "tool",
626
+ "tool_call_id": msg.get("tool_call_id"),
627
+ "content": summary
628
+ })
629
+ continue
630
+ elif role == "assistant" and msg.get("tool_calls"):
631
+ # Keep tool calls metadata so the model's message-tool call mapping DAG doesn't break
632
+ pass
633
+
634
+ # Keep general messages, but truncate if they are too long
635
+ if len(content) > 2000:
636
+ msg = dict(msg)
637
+ msg["content"] = content[:2000] + "\n[Content truncated for compaction]"
638
+ compacted.append(msg)
639
+
640
+ compacted.extend(recent_messages)
641
+ log_activity(f"[Auto-Compaction] Compressed message history from {len(messages)} down to {len(compacted)}")
642
+ return compacted
643
+
644
+
645
  @app.post("/v1/chat/completions")
646
  async def chat_completions(request: Request, authorization: str = Header(None)):
647
  auth(authorization)
 
681
  else:
682
  final_messages = list(messages)
683
 
684
+ # Perform auto-compaction before executing agent loops
685
+ final_messages = compact_history(final_messages)
686
+
687
  # Save the user's message to DB
688
  user_msg = next((m for m in reversed(messages) if m.get("role") == "user"), None)
689
  if user_msg:
 
719
  async with completions_semaphore:
720
  try:
721
  for round_num in range(MAX_TOOL_ROUNDS + 1):
722
+ # Perform auto-compaction before calling NIM API
723
+ final_messages = compact_history(final_messages)
724
  kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
725
  if is_agentic:
726
  kwargs["tools"] = TOOLS
 
832
  yield make_chunk(request_id, requested_model, "\n")
833
 
834
  # Execute the tool
835
+ result = await execute_tool(func_name, repaired_args)
836
 
837
  # Append teaching note if repaired
838
  if repair_notes: