Added timeout and limit to script entries

This commit is contained in:
Niels Geens
2024-11-13 17:48:39 +01:00
parent 65478e0b17
commit 834ce12e8f
21 changed files with 262 additions and 141 deletions

View File

@@ -17,14 +17,19 @@ class BackgroundEntry(Entry):
_exit_code: Exit code when process completes
"""
def __init__(self, script: str, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
script: str,
):
"""
Initialize a new background entry.
Args:
script: The script/command to execute
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
script: The script/command to execute
"""
super().__init__(id, timestamp)
self._script = script

View File

@@ -8,15 +8,21 @@ class ParseErrorEntry(Entry):
Entry type for parse and validation errors.
"""
def __init__(self, content: str, error: str, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
content: str,
error: str,
):
"""
Initialize a new parse error entry.
Args:
content: Original content that failed to parse
error: Error message describing the failure
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
content: Original content that failed to parse
error: Error message describing the failure
"""
super().__init__(id, timestamp)
self._content = content

View File

@@ -13,14 +13,19 @@ class ReadEntry(Entry):
_read: Whether content has been read
"""
def __init__(self, io_buffer, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
io_buffer,
):
"""
Initialize a new read entry.
Args:
io_buffer: Buffer to use for IO operations
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
io_buffer: Buffer to use for IO operations
"""
super().__init__(id, timestamp)
self._content = ""

View File

@@ -11,14 +11,19 @@ class ReasoningEntry(Entry):
_content: The reasoning text
"""
def __init__(self, content: str, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
content: str,
):
"""
Initialize a new reasoning entry.
Args:
content: The reasoning text
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
content: The reasoning text
"""
super().__init__(id, timestamp)
self._content = content

View File

@@ -8,22 +8,46 @@ from .entry import Entry
class RepeatEntry(Entry):
"""
Entry type for scripts that are executed on every update.
Attributes:
script: The script/command to execute
timeout: Maximum time to wait for script execution
limit: Maximum number of characters to capture from stdout/stderr
stdout: Captured standard output from script execution
stderr: Captured standard error from script execution
exit_code: Process exit code after completion
_executed: Whether the script has been executed
"""
default_timeout = 1
default_limit = 1024
def __init__(self, script: str, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
script: str,
timeout: Optional[float],
limit: Optional[int],
):
"""
Initialize a new repeat entry.
Args:
script: The script/command to execute
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
script: The script/command to execute
timeout: Maximum time to wait for script execution
limit: Maximum number of characters to capture from stdout/stderr
"""
super().__init__(id, timestamp)
self._script = script
self._timeout = timeout or self.default_timeout
self._limit = limit or self.default_limit
self._stdout = ""
self._stderr = ""
self._exit_code: Optional[int] = None
self._timed_out = False
@property
def script(self) -> str:
@@ -50,15 +74,20 @@ class RepeatEntry(Entry):
Execute the script and update the output.
Captures stdout, stderr and exit code from each execution.
"""
process = subprocess.run(
self._script,
shell=True,
capture_output=True,
text=True
)
self._stdout = process.stdout
self._stderr = process.stderr
self._exit_code = process.returncode
try:
process = subprocess.run(
self._script,
timeout=self._timeout,
shell=True,
capture_output=True,
text=True
)
self._stdout = process.stdout
self._stderr = process.stderr
self._exit_code = process.returncode
except subprocess.TimeoutExpired as e:
self._timed_out = True
def generate_context(self) -> ET.Element:
"""
@@ -69,12 +98,20 @@ class RepeatEntry(Entry):
"""
element = ET.Element("repeat", {
"id": self.id,
"exit_code": str(self._exit_code) if self._exit_code is not None else "-1"
})
element.text = self.script
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self._stdout
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self._stderr
if self._exit_code is not None:
element.set("exit_code", str(self._exit_code))
if self._timed_out:
element.set("timed_out", "true")
else:
if len(self._stdout) > self._limit:
element.set("stdout_truncated", "true")
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self._stdout[:self._limit]
if len(self._stderr) > self._limit:
element.set("stderr_truncated", "true")
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self._stderr[:self._limit]
return element

View File

@@ -54,15 +54,15 @@ class ResponseParser:
try:
root = ET.fromstring(xml.strip())
except ET.ParseError as e:
return ParseErrorEntry(xml, f"Invalid XML: {str(e)}", entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, f"Invalid XML: {str(e)}")
try:
if root.tag == 'delete':
target_id = root.get('id')
if not target_id:
return ParseErrorEntry(xml, "Delete command missing required 'id' attribute", entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, "Delete command missing required 'id' attribute")
if len(root.attrib) > 1:
return ParseErrorEntry(xml, "Delete command should only have 'id' attribute", entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, "Delete command should only have 'id' attribute")
return DeleteCommand(target_id)
elif root.tag == 'stop':
@@ -70,34 +70,46 @@ class ResponseParser:
elif root.tag == 'background':
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
return ParseErrorEntry(xml, "Background entry requires (only) script content", entry_id, timestamp)
return BackgroundEntry(root.text, entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, "Background entry requires (only) script content")
return BackgroundEntry(entry_id, timestamp, root.text)
elif root.tag == 'repeat':
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
return ParseErrorEntry(xml, "Repeat entry requires (only) script content", entry_id, timestamp)
return RepeatEntry(root.text, entry_id, timestamp)
if len(root) != 0 or root.text is None or root.text.strip() == '':
return ParseErrorEntry(entry_id, timestamp, xml, "Repeat entry requires (only) script content")
if len(root.attrib) > 2 or any(k not in ('timeout', 'limit') for k in root.attrib):
return ParseErrorEntry(entry_id, timestamp, xml, "Repeat entry only accepts 'timeout' and 'limit' attributes")
timeout = root.get('timeout')
limit = root.get('limit')
timeout = float(timeout) if timeout is not None else None
limit = int(limit) if limit is not None else None
return RepeatEntry(entry_id, timestamp, root.text, timeout, limit)
elif root.tag == 'single':
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
return ParseErrorEntry(xml, "Single entry requires (only) script content", entry_id, timestamp)
return SingleEntry(root.text, entry_id, timestamp)
if len(root) != 0 or root.text is None or root.text.strip() == '':
return ParseErrorEntry(entry_id, timestamp, xml, "Single entry requires (only) script content")
if len(root.attrib) > 2 or any(k not in ('timeout', 'limit') for k in root.attrib):
return ParseErrorEntry(entry_id, timestamp, xml, "Single entry only accepts 'timeout' and 'limit' attributes")
timeout = root.get('timeout')
limit = root.get('limit')
timeout = float(timeout) if timeout is not None else None
limit = int(limit) if limit is not None else None
return SingleEntry(entry_id, timestamp, root.text, timeout, limit)
elif root.tag == 'reasoning':
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
return ParseErrorEntry(xml, "Reasoning entry requires (only) text content", entry_id, timestamp)
return ReasoningEntry(root.text, entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, "Reasoning entry requires (only) text content")
return ReasoningEntry(entry_id, timestamp, root.text)
elif root.tag == 'read_stdin':
return ReadEntry(self.io_buffer, entry_id, timestamp)
return ReadEntry(entry_id, timestamp, self.io_buffer)
elif root.tag == 'write_stdout':
if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
return ParseErrorEntry(xml, "Write stdout entry requires (only) text content", entry_id, timestamp)
return WriteEntry(root.text, self.io_buffer, entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, "Write stdout entry requires (only) text content")
return WriteEntry(entry_id, timestamp, root.text, self.io_buffer)
else:
return ParseErrorEntry(xml, f"Unknown root element: {root.tag}", entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, f"Unknown root element: {root.tag}")
except Exception as e:
return ParseErrorEntry(xml, f"Error parsing response: {str(e)}", entry_id, timestamp)
return ParseErrorEntry(entry_id, timestamp, xml, f"Error parsing response: {str(e)}")

View File

@@ -11,27 +11,44 @@ class SingleEntry(Entry):
Attributes:
script: The script/command to execute
timeout: Maximum time to wait for script execution
limit: Maximum number of characters to capture from stdout/stderr
stdout: Captured standard output from script execution
stderr: Captured standard error from script execution
exit_code: Process exit code after completion
_executed: Whether the script has been executed
"""
default_timeout = 1
default_limit = 1024
def __init__(self, script: str, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
script: str,
timeout: Optional[float],
limit: Optional[int],
):
"""
Initialize a new single shot entry.
Args:
script: The script/command to execute
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
script: The script/command to execute
timeout: Maximum time to wait for script execution
limit: Maximum number of characters to capture from stdout/stderr
"""
super().__init__(id, timestamp)
self._script = script
self._timeout = timeout or self.default_timeout
self._limit = limit or self.default_limit
self._stdout = ""
self._stderr = ""
self._exit_code: Optional[int] = None
self._executed = False
self._timed_out = False
@property
def script(self) -> str:
@@ -60,16 +77,20 @@ class SingleEntry(Entry):
"""
if self._executed:
return
process = subprocess.run(
self.script,
shell=True,
capture_output=True,
text=True
)
self._stdout = process.stdout
self._stderr = process.stderr
self._exit_code = process.returncode
self._executed = True
try:
process = subprocess.run(
self._script,
timeout=self._timeout,
shell=True,
capture_output=True,
text=True
)
self._stdout = process.stdout
self._stderr = process.stderr
self._exit_code = process.returncode
except subprocess.TimeoutExpired as e:
self._timed_out = True
def generate_context(self) -> ET.Element:
"""
@@ -82,10 +103,16 @@ class SingleEntry(Entry):
"id": self.id,
})
element.text = self.script
if self._executed:
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self._stdout
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self._stderr
if self._timed_out:
element.set("timed_out", "true")
elif self._executed:
element.set("exit_code", str(self._exit_code))
if len(self._stdout) > self._limit:
element.set("stdout_truncated", "true")
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self._stdout[:self._limit]
if len(self._stderr) > self._limit:
element.set("stderr_truncated", "true")
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self._stderr[:self._limit]
return element

View File

@@ -2,6 +2,7 @@ from datetime import datetime
import xml.etree.ElementTree as ET
from .entry import Entry
from .io_buffer import IOBuffer
class WriteEntry(Entry):
"""
@@ -13,15 +14,21 @@ class WriteEntry(Entry):
_written: Whether the content has been written
"""
def __init__(self, content: str, io_buffer, id: str, timestamp: datetime):
def __init__(
self,
id: str,
timestamp: datetime,
content: str,
io_buffer: IOBuffer,
):
"""
Initialize a new write entry.
Args:
content: Text to write to stdout
io_buffer: Buffer to use for IO operations
id: Unique identifier for this entry
timestamp: Creation timestamp for this entry
content: Text to write to stdout
io_buffer: Buffer to use for IO operations
"""
super().__init__(id, timestamp)
self.content = content