diff --git a/docs/supervisor-usage.md b/docs/supervisor-usage.md index cef3a44..c5778bf 100644 --- a/docs/supervisor-usage.md +++ b/docs/supervisor-usage.md @@ -138,7 +138,7 @@ python -m supervisor.supervisor [OPTIONS] | Argument | Short | Required | Default | Description | |----------|-------|----------|---------|-------------| | `--config-file` | `-f` | Yes | - | Path to task configuration YAML | -| `--duration` | `-d` | No | 60 | Duration to run (minutes, pauses during non-working hours) | +| `--duration` | `-d` | No | 60 | Duration to run (minutes, pauses during non-working hours if enabled) | | `--supervisor-model` | `-m` | No | `openai/o4-mini` | Model for supervisor LLM | | `--resume-dir` | - | No | - | Resume from existing session directory | | `--verbose` | `-v` | No | False | Enable verbose logging (DEBUG level) | @@ -147,9 +147,6 @@ python -m supervisor.supervisor [OPTIONS] | `--skip-todos` | - | No | False | Skip initial TODO generation step | | `--use-prompt-generation` | - | No | False | Use LLM to generate custom system prompts instead of routing | | `--finish-on-submit` | - | No | False | Finish session when a vulnerability is submitted (instead of continuing until duration expires) | -| `--working-hours-start` | - | No | 9 | Working hours start time (24-hour format, e.g., 9 for 9 AM) | -| `--working-hours-end` | - | No | 17 | Working hours end time (24-hour format, e.g., 17 for 5 PM) | -| `--working-hours-timezone` | - | No | `US/Pacific` | Timezone for working hours (e.g., `US/Pacific`, `UTC`, `US/Eastern`) | ## Modes @@ -250,35 +247,56 @@ python -m supervisor.supervisor \ ``` ### Working Hours Configuration + +Working hours are configured in the YAML config file, not via CLI flags. See the [Working Hours](#working-hours) section for details. + +```yaml +# In your config YAML file +working_hours: + enabled: true + start_hour: 9 + end_hour: 17 + timezone: "US/Eastern" +``` + ```bash -# Run supervisor only during business hours in Eastern timezone +# Run supervisor with working hours enabled in config python -m supervisor.supervisor \ --config-file ../configs/stanford/level1.yaml \ - --duration 480 \ - --working-hours-start 9 \ - --working-hours-end 17 \ - --working-hours-timezone "US/Eastern" + --duration 480 ``` ## Advanced Features ### Working Hours -The supervisor supports **automatic sleep during non-working hours**, ensuring operations only occur during specified business hours. The session duration automatically adjusts to account for sleep time. +The supervisor supports **automatic sleep during non-working hours**, ensuring operations only occur during specified business hours. This feature is **disabled by default** and must be explicitly enabled in your task configuration YAML. **Features:** +- Disabled by default (runs 24/7 unless configured) - Automatically pauses supervisor when outside working hours - Session duration extended by sleep time (e.g., 60 minute duration = 60 minutes of actual work) - Timezone-aware scheduling - Working hours status shown in heartbeat file and logs -- Configurable via command-line arguments +- Always disabled in benchmark mode (even if configured) -**Configuration:** -- `--working-hours-start` - Start hour in 24-hour format (default: 9 for 9 AM) -- `--working-hours-end` - End hour in 24-hour format (default: 17 for 5 PM) -- `--working-hours-timezone` - Timezone string (default: `US/Pacific`) +**Configuration (in task YAML file):** -**Example behavior:** +```yaml +working_hours: + enabled: true # Required to enable (default: disabled) + start_hour: 9 # 24-hour format (default: 9) + end_hour: 17 # 24-hour format (default: 17) + timezone: "US/Pacific" # Timezone string (default: US/Pacific) +``` + +**Behavior:** +- No `working_hours` block in config: runs continuously (24/7) +- `working_hours.enabled: false`: runs continuously +- `working_hours.enabled: true`: pauses outside configured hours +- Benchmark mode (`--benchmark-mode`): always runs continuously, ignores config + +**Example behavior (when enabled):** - Start session at 4:00 PM with 120 minute duration - Working hours: 9 AM - 5 PM - Supervisor runs for 60 minutes (4:00 PM - 5:00 PM) @@ -595,6 +613,13 @@ targets: # Files referenced in configuration (relative paths resolved from config file directory) filepath: "./additional_context.txt" # Optional +# Working hours scheduling (optional, disabled by default) +working_hours: + enabled: true # Set to true to enable + start_hour: 9 # 24-hour format + end_hour: 17 # 24-hour format + timezone: "US/Pacific" # Timezone string + # Submission configuration (for benchmark mode) submission_config: type: "ctf" # or "vulnerability" @@ -619,6 +644,7 @@ testing_credentials: - At minimum one of: `targets`, `description`, or custom task-specific fields **Optional Fields:** +- `working_hours` - Working hours scheduling (disabled by default) - `submission_config` - Required for benchmark mode - `filepath` - Additional context file - Any custom fields needed for your specific use case diff --git a/pyproject.toml b/pyproject.toml index 6c51a7a..8cbe14a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "python-dotenv>=1.0.0", "openai>=2.3.0", "tiktoken>=0.12.0", + "pydantic>=2.12.0", ] [build-system] diff --git a/supervisor/config.py b/supervisor/config.py new file mode 100644 index 0000000..43ba6de --- /dev/null +++ b/supervisor/config.py @@ -0,0 +1,16 @@ +from pydantic import BaseModel, Field, model_validator + + +class WorkingHoursConfig(BaseModel, frozen=True): + """Configuration for working hours scheduling.""" + + enabled: bool = False + start_hour: int = Field(default=9, ge=0, le=23) + end_hour: int = Field(default=17, ge=0, le=23) + timezone: str = "US/Pacific" + + @model_validator(mode="after") + def validate_hours(self) -> "WorkingHoursConfig": + if self.start_hour >= self.end_hour: + raise ValueError(f"start_hour ({self.start_hour}) must be < end_hour ({self.end_hour})") + return self diff --git a/supervisor/orchestration/orchestrator.py b/supervisor/orchestration/orchestrator.py index c23b1a9..ce24290 100644 --- a/supervisor/orchestration/orchestrator.py +++ b/supervisor/orchestration/orchestrator.py @@ -17,7 +17,8 @@ from ..prompts.summarization_prompt import get_summarization_prompt from ..prompts.supervisor_prompt import SupervisorPrompt from ..context_manager import ContextManager -from ..working_hours import WorkingHoursManager +from supervisor.working_hours import WorkingHoursManager +from supervisor.config import WorkingHoursConfig from .instance_manager import InstanceManager from .log_reader import LogReader @@ -29,9 +30,9 @@ class SupervisorOrchestrator: def __init__(self, config: Dict[str, Any], session_dir: Path, supervisor_model: str = "o3", duration_minutes: int = 60, verbose: bool = False, codex_binary: str = "./target/release/codex", benchmark_mode: bool = False, skip_todos: bool = False, use_prompt_generation: bool = False, - working_hours_start: int = 9, working_hours_end: int = 17, working_hours_timezone: str = "US/Pacific", + working_hours_config: Optional[WorkingHoursConfig] = None, finish_on_submit: bool = False): - + self.config = config self.session_dir = session_dir self.supervisor_model = supervisor_model @@ -42,13 +43,16 @@ def __init__(self, config: Dict[str, Any], session_dir: Path, supervisor_model: self.skip_todos = skip_todos self.use_prompt_generation = use_prompt_generation self.finish_on_submit = finish_on_submit - - # Initialize working hours manager - self.working_hours = WorkingHoursManager( - start_hour=working_hours_start, - end_hour=working_hours_end, - timezone_str=working_hours_timezone - ) + + # Working hours scheduling (disabled by default, always disabled in benchmark mode) + self.working_hours: Optional[WorkingHoursManager] = None + wh = working_hours_config or WorkingHoursConfig() + if wh.enabled and not benchmark_mode: + self.working_hours = WorkingHoursManager( + start_hour=wh.start_hour, + end_hour=wh.end_hour, + timezone_str=wh.timezone, + ) self.instance_manager = InstanceManager(session_dir, codex_binary, use_prompt_generation=use_prompt_generation) self.log_reader = LogReader(session_dir, self.instance_manager) @@ -129,14 +133,17 @@ async def run_loop(self): logging.info(f"đŸŽ¯ Supervisor starting {self.duration_minutes}min session") logging.info(f"📅 Session will end at: {end_time.strftime('%Y-%m-%d %H:%M:%S UTC')}") - - # Log working hours information - working_status = self.working_hours.get_status_info() - logging.info(f"🕐 Working hours: {working_status['working_hours']}") - if working_status['in_working_hours']: - logging.info(f"✅ Currently in working hours, time until end: {working_status['time_until_end']}") + + # Log working hours status + if self.working_hours: + status = self.working_hours.get_status_info() + logging.info(f"Working hours enabled: {status['working_hours']}") + if status['in_working_hours']: + logging.info(f"Currently in working hours, ends in: {status['time_until_end']}") + else: + logging.info(f"Outside working hours, resumes at: {status['next_working_time']}") else: - logging.info(f"⏰ Currently outside working hours, next working time: {working_status['next_working_time']}") + logging.info("Working hours disabled, running continuously") await self._save_session_metadata(start_time, end_time) @@ -146,16 +153,15 @@ async def run_loop(self): try: iteration += 1 logging.info(f"🔄 Supervisor iteration {iteration}") - - # Check working hours and sleep if necessary - sleep_duration, wake_time = await self.working_hours.wait_for_working_hours() - if sleep_duration.total_seconds() > 0: - self.sleep_time_outside_hours += sleep_duration - logging.info(f"⏱ī¸ Total sleep time outside working hours: {self._format_duration(self.sleep_time_outside_hours)}") - - # Update heartbeat after waking up - await self._update_heartbeat(iteration, start_time, sleeping=False) - + + # Sleep until working hours if configured + if self.working_hours: + sleep_duration, _ = await self.working_hours.wait_for_working_hours() + if sleep_duration.total_seconds() > 0: + self.sleep_time_outside_hours += sleep_duration + logging.info(f"Slept {self._format_duration(sleep_duration)} outside working hours") + await self._update_heartbeat(iteration, start_time, sleeping=False) + await self._update_heartbeat(iteration, start_time, sleeping=False) user_message = await self._generate_instance_update_message() @@ -416,8 +422,17 @@ async def _get_supervisor_response(self, instance_responses: Dict[str, str] = No async def _update_heartbeat(self, iteration: int, start_time: datetime, sleeping: bool = False): """Update supervisor heartbeat file.""" - working_status = self.working_hours.get_status_info() - + if self.working_hours: + status = self.working_hours.get_status_info() + working_hours_info = { + "enabled": True, + "config": status['working_hours'], + "in_working_hours": status['in_working_hours'], + "total_sleep_time": self._format_duration(self.sleep_time_outside_hours) + } + else: + working_hours_info = {"enabled": False} + heartbeat = { "supervisor_pid": os.getpid(), "session_dir": str(self.session_dir), @@ -426,11 +441,7 @@ async def _update_heartbeat(self, iteration: int, start_time: datetime, sleeping "start_time": start_time.isoformat(), "active_instances": len([i for i in self.instance_manager.instances.values() if i["status"] == "running"]), "status": "sleeping" if sleeping else "running", - "working_hours": { - "config": working_status['working_hours'], - "in_working_hours": working_status['in_working_hours'], - "total_sleep_time": self._format_duration(self.sleep_time_outside_hours) - } + "working_hours": working_hours_info } try: diff --git a/supervisor/supervisor.py b/supervisor/supervisor.py index a4480b6..65ce25a 100644 --- a/supervisor/supervisor.py +++ b/supervisor/supervisor.py @@ -13,8 +13,9 @@ from dotenv import load_dotenv load_dotenv() -from .orchestration import SupervisorOrchestrator -from .todo_generator import TodoGenerator +from supervisor.orchestration import SupervisorOrchestrator +from supervisor.todo_generator import TodoGenerator +from supervisor.config import WorkingHoursConfig def setup_logging(session_dir: Path, verbose: bool = False): """Setup logging for the supervisor.""" @@ -67,12 +68,6 @@ async def main(): help='Skip the initial TODO generation step') parser.add_argument('--use-prompt-generation', action='store_true', help='Use LLM to generate custom system prompts instead of routing to predefined modes') - parser.add_argument('--working-hours-start', type=int, default=9, - help='Working hours start time (24-hour format, default: 9 for 9 AM)') - parser.add_argument('--working-hours-end', type=int, default=17, - help='Working hours end time (24-hour format, default: 17 for 5 PM)') - parser.add_argument('--working-hours-timezone', default='US/Pacific', - help='Timezone for working hours (default: US/Pacific)') parser.add_argument('--finish-on-submit', action='store_true', help='Finish session when a vulnerability is submitted (instead of continuing until duration expires)') @@ -135,7 +130,7 @@ async def main(): if args.finish_on_submit: print("⏚ī¸ FINISH ON SUBMIT MODE ENABLED - Session will end after first submission") - + codex_binary_path = Path(args.codex_binary).resolve() if not codex_binary_path.exists(): print(f"❌ Codex binary not found: {codex_binary_path}") @@ -184,9 +179,7 @@ def count_all_todos(todos): benchmark_mode=args.benchmark_mode, skip_todos=args.skip_todos, use_prompt_generation=args.use_prompt_generation, - working_hours_start=args.working_hours_start, - working_hours_end=args.working_hours_end, - working_hours_timezone=args.working_hours_timezone, + working_hours_config=WorkingHoursConfig.model_validate(config.get('working_hours', {})), finish_on_submit=args.finish_on_submit ) diff --git a/uv.lock b/uv.lock index 0575a6d..b976abf 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.12" [[package]] @@ -93,6 +93,7 @@ dependencies = [ { name = "aiofiles" }, { name = "openai" }, { name = "psutil" }, + { name = "pydantic" }, { name = "python-dotenv" }, { name = "pytz" }, { name = "pyyaml" }, @@ -108,12 +109,13 @@ dev = [ [package.metadata] requires-dist = [ { name = "aiofiles", specifier = ">=24.0.0" }, - { name = "openai", specifier = ">=1.0.0" }, + { name = "openai", specifier = ">=2.3.0" }, { name = "psutil", specifier = ">=5.9.0" }, + { name = "pydantic", specifier = ">=2.12.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, { name = "pytz", specifier = ">=2023.3" }, { name = "pyyaml", specifier = ">=6.0" }, - { name = "tiktoken", specifier = ">=0.5.0" }, + { name = "tiktoken", specifier = ">=0.12.0" }, ] [package.metadata.requires-dev]