diff --git a/.egg-state/checks/implement-results.json b/.egg-state/checks/implement-results.json new file mode 100644 index 0000000000..b165a18ac8 --- /dev/null +++ b/.egg-state/checks/implement-results.json @@ -0,0 +1,20 @@ +{ + "all_passed": false, + "checks": [ + { + "name": "ruff-check", + "passed": true, + "output": "All checks passed!" + }, + { + "name": "ruff-format", + "passed": false, + "output": "29 files would be reformatted, 311 files already formatted. Formatting issues in gateway/, orchestrator/, shared/, tests/ directories." + }, + { + "name": "pytest", + "passed": false, + "output": "4123 passed, 3 failed, 81 skipped, 2 errors in 39.39s. Failures: 3 tests in gateway/tests/test_session_manager.py (TestSessionEndCheckpointCapture - _capture_and_cleanup_session not called). Errors: 2 tests in gateway/tests/test_worktree_manager.py (TestWorktreeManagerDockerGitDir - git init failed in sandbox environment)." + } + ] +} diff --git a/.egg-state/contracts/530.json b/.egg-state/contracts/530.json new file mode 100644 index 0000000000..e5c9178979 --- /dev/null +++ b/.egg-state/contracts/530.json @@ -0,0 +1,1031 @@ +{ + "schemaVersion": "1.0", + "issue": { + "number": 530, + "title": "Ensure checkpoints are captured if an agent doesn't push to github", + "url": "https://github.com/jwbron/egg/issues/530" + }, + "pipeline_id": null, + "current_phase": "implement", + "acceptance_criteria": [ + { + "id": "ac-1", + "description": "[TASK-1-1] Enums defined with all values from spec, unit tests pass", + "verified": true + }, + { + "id": "ac-2", + "description": "[TASK-1-2] Model validates correctly, supports both commit and session-end checkpoints", + "verified": true + }, + { + "id": "ac-3", + "description": "[TASK-1-3] Summary captures all fields needed for index lookups", + "verified": true + }, + { + "id": "ac-4", + "description": "[TASK-1-4] Index model supports O(1) lookups for all dimensions", + "verified": true + }, + { + "id": "ac-5", + "description": "[TASK-1-5] Tests cover validation, serialization, and edge cases", + "verified": true + }, + { + "id": "ac-6", + "description": "[TASK-2-1] Generates unique IDs without requiring commit_sha", + "verified": true + }, + { + "id": "ac-7", + "description": "[TASK-2-2] Writes to v2 directory structure, atomic file operations", + "verified": true + }, + { + "id": "ac-8", + "description": "[TASK-2-3] Correctly loads v2 checkpoints and index", + "verified": true + }, + { + "id": "ac-9", + "description": "[TASK-2-4] Updates all secondary indices atomically", + "verified": true + }, + { + "id": "ac-10", + "description": "[TASK-2-5] O(1) lookups using secondary indices", + "verified": true + }, + { + "id": "ac-11", + "description": "[TASK-2-6] Tests cover save/load, index updates, and concurrent access patterns", + "verified": true + }, + { + "id": "ac-12", + "description": "[TASK-3-1] Fields persisted to disk, backward compatible with existing sessions", + "verified": true + }, + { + "id": "ac-13", + "description": "[TASK-3-2] Parameters stored in Session, passed to gateway API", + "verified": true + }, + { + "id": "ac-14", + "description": "[TASK-3-3] New fields round-trip correctly through persistence", + "verified": true + }, + { + "id": "ac-15", + "description": "[TASK-3-4] Tests cover new fields, backward compatibility", + "verified": true + }, + { + "id": "ac-16", + "description": "[TASK-4-1] Longer messages preserved, existing tests pass", + "verified": true + }, + { + "id": "ac-17", + "description": "[TASK-4-2] Tool parameters less aggressively truncated", + "verified": true + }, + { + "id": "ac-18", + "description": "[TASK-4-3] Tool results more complete", + "verified": true + }, + { + "id": "ac-19", + "description": "[TASK-4-4] Larger transcripts accepted", + "verified": true + }, + { + "id": "ac-20", + "description": "[TASK-4-5] Comments reflect new limits and rationale", + "verified": true + }, + { + "id": "ac-21", + "description": "[TASK-5-1] Captures checkpoint with trigger_type=SESSION_END, uses v2 schema", + "verified": true + }, + { + "id": "ac-22", + "description": "[TASK-5-2] Stores to egg/checkpoints/v2 branch with proper index updates", + "verified": true + }, + { + "id": "ac-23", + "description": "[TASK-5-3] Buffer not deleted until capture completes or timeout", + "verified": true + }, + { + "id": "ac-24", + "description": "[TASK-5-4] Transcript marked with truncated_reason when applicable", + "verified": true + }, + { + "id": "ac-25", + "description": "[TASK-5-5] Existing push-triggered checkpoints use v2 format", + "verified": true + }, + { + "id": "ac-26", + "description": "[TASK-6-1] Graceful deletions produce checkpoints", + "verified": true + }, + { + "id": "ac-27", + "description": "[TASK-6-2] Container-based deletions produce checkpoints", + "verified": true + }, + { + "id": "ac-28", + "description": "[TASK-6-3] TTL expirations produce checkpoints", + "verified": true + }, + { + "id": "ac-29", + "description": "[TASK-6-4] Session metadata available for orphan checkpoints", + "verified": true + }, + { + "id": "ac-30", + "description": "[TASK-6-5] Crashed containers produce checkpoints", + "verified": true + }, + { + "id": "ac-31", + "description": "[TASK-6-6] Buffer preserved until capture finishes or 30s timeout", + "verified": true + }, + { + "id": "ac-32", + "description": "[TASK-7-1] New checkpoints written to v2 branch", + "verified": true + }, + { + "id": "ac-33", + "description": "[TASK-7-2] Push checkpoints use v2 models", + "verified": true + }, + { + "id": "ac-34", + "description": "[TASK-7-3] Multi-commit pushes produce v2 checkpoints", + "verified": true + }, + { + "id": "ac-35", + "description": "[TASK-7-4] No v1 code paths remain active", + "verified": true + }, + { + "id": "ac-36", + "description": "[TASK-8-1] Test covers graceful deletion, TTL expiration, and crash scenarios", + "verified": true + }, + { + "id": "ac-37", + "description": "[TASK-8-2] Push-triggered checkpoints work with v2", + "verified": true + }, + { + "id": "ac-38", + "description": "[TASK-8-3] Queries by session, issue, PR, status all work", + "verified": true + }, + { + "id": "ac-39", + "description": "[TASK-8-4] Code is well-documented for future maintainers", + "verified": true + }, + { + "id": "ac-40", + "description": "[TASK-8-5] Full test suite passes, manual smoke test succeeds", + "verified": true + } + ], + "phases": [ + { + "id": "phase-1", + "name": "Core Contracts & Models", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-1-1", + "description": "Add TriggerType, SessionStatus, and AgentType enums to checkpoints.py", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Enums defined with all values from spec, unit tests pass", + "files_affected": [ + "shared/egg_contracts/checkpoints.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-1-2", + "description": "Add CheckpointV2 model with optional commit_sha, required trigger_type and session_id", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Model validates correctly, supports both commit and session-end checkpoints", + "files_affected": [ + "shared/egg_contracts/checkpoints.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-1-3", + "description": "Add CheckpointSummaryV2 model with all queryable fields", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Summary captures all fields needed for index lookups", + "files_affected": [ + "shared/egg_contracts/checkpoints.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-1-4", + "description": "Add CheckpointIndexV2 model with secondary indices", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Index model supports O(1) lookups for all dimensions", + "files_affected": [ + "shared/egg_contracts/checkpoints.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-1-5", + "description": "Add unit tests for v2 models", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Tests cover validation, serialization, and edge cases", + "files_affected": [ + "shared/egg_contracts/tests/test_checkpoints.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-2", + "name": "Checkpoint Loader v2 Functions", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-2-1", + "description": "Add generate_checkpoint_id_v2() function for session-end checkpoints", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Generates unique IDs without requiring commit_sha", + "files_affected": [ + "shared/egg_contracts/checkpoint_loader.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-2-2", + "description": "Add save_checkpoint_v2() with atomic write support", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Writes to v2 directory structure, atomic file operations", + "files_affected": [ + "shared/egg_contracts/checkpoint_loader.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-2-3", + "description": "Add load_checkpoint_v2() and load_checkpoint_index_v2() functions", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Correctly loads v2 checkpoints and index", + "files_affected": [ + "shared/egg_contracts/checkpoint_loader.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-2-4", + "description": "Add add_checkpoint_to_index_v2() with multi-index updates", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Updates all secondary indices atomically", + "files_affected": [ + "shared/egg_contracts/checkpoint_loader.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-2-5", + "description": "Add v2-specific lookup helpers (get_by_session, get_by_trigger, get_by_status)", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "O(1) lookups using secondary indices", + "files_affected": [ + "shared/egg_contracts/checkpoint_loader.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-2-6", + "description": "Add unit tests for v2 loader functions", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Tests cover save/load, index updates, and concurrent access patterns", + "files_affected": [ + "shared/egg_contracts/tests/test_checkpoint_loader.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-3", + "name": "Session Metadata Enhancement", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-3-1", + "description": "Add issue_number and pr_number optional fields to Session dataclass", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Fields persisted to disk, backward compatible with existing sessions", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-3-2", + "description": "Update register_session() to accept optional issue_number and pr_number", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Parameters stored in Session, passed to gateway API", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-3-3", + "description": "Update Session persistence methods for new fields", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "New fields round-trip correctly through persistence", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-3-4", + "description": "Add unit tests for session metadata", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Tests cover new fields, backward compatibility", + "files_affected": [ + "gateway/tests/test_session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-4", + "name": "Transcript Size Limit Increase", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-4-1", + "description": "Increase max_content_length from 10,000 to 25,000 characters", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Longer messages preserved, existing tests pass", + "files_affected": [ + "shared/egg_contracts/transcript_extractor.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-4-2", + "description": "Increase max_param_length from 1,000 to 2,500 characters", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Tool parameters less aggressively truncated", + "files_affected": [ + "shared/egg_contracts/transcript_extractor.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-4-3", + "description": "Increase max_result_length from 500 to 1,500 characters", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Tool results more complete", + "files_affected": [ + "shared/egg_contracts/transcript_extractor.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-4-4", + "description": "Increase MAX_TRANSCRIPT_SIZE from 1MB to 3MB", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Larger transcripts accepted", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-4-5", + "description": "Update constants documentation", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Comments reflect new limits and rationale", + "files_affected": [ + "shared/egg_contracts/transcript_extractor.py", + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-5", + "name": "Session-End Checkpoint Capture", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-5-1", + "description": "Add capture_session_end_checkpoint() function", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Captures checkpoint with trigger_type=SESSION_END, uses v2 schema", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-5-2", + "description": "Add store_checkpoint_v2() method for v2 branch", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Stores to egg/checkpoints/v2 branch with proper index updates", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-5-3", + "description": "Implement async capture with buffer preservation (30s timeout)", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Buffer not deleted until capture completes or timeout", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-5-4", + "description": "Add truncated_reason field handling for crashed sessions", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Transcript marked with truncated_reason when applicable", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-5-5", + "description": "Update capture_checkpoint() to produce v2 checkpoints", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Existing push-triggered checkpoints use v2 format", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-6", + "name": "Integration with Session Deletion Paths", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-6-1", + "description": "Update delete_session() to capture checkpoint with status=COMPLETED", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Graceful deletions produce checkpoints", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-6-2", + "description": "Update delete_session_by_container() to capture checkpoint", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Container-based deletions produce checkpoints", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-6-3", + "description": "Update prune_expired_sessions() to capture checkpoints with status=EXPIRED", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "TTL expirations produce checkpoints", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-6-4", + "description": "Add session lookup before worktree cleanup", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Session metadata available for orphan checkpoints", + "files_affected": [ + "gateway/worktree_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-6-5", + "description": "Update cleanup_orphaned_worktrees() to capture checkpoint with status=FAILED", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Crashed containers produce checkpoints", + "files_affected": [ + "gateway/worktree_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-6-6", + "description": "Refactor _cleanup_transcript_buffer() to await checkpoint completion", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Buffer preserved until capture finishes or 30s timeout", + "files_affected": [ + "gateway/session_manager.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-7", + "name": "Update Push-Triggered Checkpoints", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-7-1", + "description": "Update CHECKPOINT_BRANCH constant to egg/checkpoints/v2", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "New checkpoints written to v2 branch", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-7-2", + "description": "Update capture_and_store_checkpoint() to use v2 capture and storage", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Push checkpoints use v2 models", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-7-3", + "description": "Update capture_and_store_checkpoints_for_push() to produce v2 checkpoints", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Multi-commit pushes produce v2 checkpoints", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-7-4", + "description": "Remove v1 imports and references", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "No v1 code paths remain active", + "files_affected": [ + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + }, + { + "id": "phase-8", + "name": "Integration Testing & Documentation", + "status": "pending", + "review_cycles": 0, + "max_cycles": 3, + "escalated": false, + "escalation_reason": null, + "tasks": [ + { + "id": "task-8-1", + "description": "Add integration tests for session-end checkpoint flow", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Test covers graceful deletion, TTL expiration, and crash scenarios", + "files_affected": [ + "gateway/tests/test_checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-8-2", + "description": "Add integration tests for commit checkpoint flow with v2", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Push-triggered checkpoints work with v2", + "files_affected": [ + "gateway/tests/test_checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-8-3", + "description": "Add integration tests for multi-dimensional index queries", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Queries by session, issue, PR, status all work", + "files_affected": [ + "gateway/tests/test_checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-8-4", + "description": "Update checkpoint docstrings and module-level documentation", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Code is well-documented for future maintainers", + "files_affected": [ + "shared/egg_contracts/checkpoints.py", + "gateway/checkpoint_handler.py" + ], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + }, + { + "id": "task-8-5", + "description": "Verify no regressions in existing gateway functionality", + "status": "pending", + "commit": null, + "checkpoint_id": null, + "notes": "", + "acceptance_criteria": "Full test suite passes, manual smoke test succeeds", + "files_affected": [], + "review_cycles": 0, + "max_cycles": 3, + "escalated": false + } + ], + "review_feedback": [] + } + ], + "decisions": [], + "workflow_owner": "jwbron", + "audit_log": [ + { + "timestamp": "2026-02-12T21:37:33Z", + "actor": "system", + "role": "reviewer", + "action": "update", + "field_path": "refine_review_cycles", + "old_value": 0, + "new_value": 1, + "reason": "Review cycle 1: approved (per-reviewer: {\"unified\":\"missing\",\"agent-design\":\"missing\"})", + "checkpoint_id": null + }, + { + "timestamp": "2026-02-12T22:26:39Z", + "actor": "jwbron", + "role": "human", + "action": "transition", + "field_path": "current_phase", + "old_value": "refine", + "new_value": "plan", + "reason": "Phase approved by human", + "checkpoint_id": null + }, + { + "timestamp": "2026-02-12T22:36:23Z", + "actor": "system", + "role": "reviewer", + "action": "update", + "field_path": "plan_review_cycles", + "old_value": 0, + "new_value": 1, + "reason": "Review cycle 1: approved (per-reviewer: {\"unified\":\"missing\",\"agent-design\":\"missing\"})", + "checkpoint_id": null + }, + { + "timestamp": "2026-02-12T22:50:02Z", + "actor": "jwbron", + "role": "human", + "action": "transition", + "field_path": "current_phase", + "old_value": "plan", + "new_value": "implement", + "reason": "Phase approved by human", + "checkpoint_id": null + } + ], + "refine_review_cycles": 1, + "refine_review_feedback": "", + "plan_review_cycles": 1, + "plan_review_feedback": "", + "pr": { + "title": "Checkpoint v2: capture all sessions with rich querying", + "description": "Implements a v2 checkpoint system that captures every agent session regardless\nof push activity. Adds session-end checkpoints (COMPLETED, EXPIRED, FAILED)\nwith multi-dimensional index for querying by session, issue, PR, commit,\nagent type, phase, and status. Resolves #530." + }, + "feedback": null, + "phase_configs": { + "refine": { + "checks": [], + "max_review_cycles": 3, + "human_review_mechanism": "ISSUE_CHECKBOX" + }, + "plan": { + "checks": [], + "max_review_cycles": 3, + "human_review_mechanism": "ISSUE_CHECKBOX" + }, + "implement": { + "checks": [], + "max_review_cycles": 3, + "human_review_mechanism": "PR_REVIEW" + } + }, + "agent_executions": [ + { + "role": "coder", + "status": "complete", + "started_at": null, + "completed_at": "2026-02-13T17:59:05.781822Z", + "commit": null, + "checkpoint_id": null, + "outputs": {}, + "error": null, + "retry_count": 0 + }, + { + "role": "tester", + "status": "pending", + "started_at": null, + "completed_at": null, + "commit": null, + "checkpoint_id": null, + "outputs": {}, + "error": null, + "retry_count": 0 + }, + { + "role": "documenter", + "status": "pending", + "started_at": null, + "completed_at": null, + "commit": null, + "checkpoint_id": null, + "outputs": {}, + "error": null, + "retry_count": 0 + }, + { + "role": "integrator", + "status": "pending", + "started_at": null, + "completed_at": null, + "commit": null, + "checkpoint_id": null, + "outputs": {}, + "error": null, + "retry_count": 0 + } + ], + "multi_agent_config": null +} diff --git a/.egg-state/drafts/530-analysis.md b/.egg-state/drafts/530-analysis.md new file mode 100644 index 0000000000..d5015420c1 --- /dev/null +++ b/.egg-state/drafts/530-analysis.md @@ -0,0 +1,312 @@ +# Analysis: Checkpoint v2 — Capture All Sessions with Rich Querying (#530) + +## Problem Statement + +The v1 checkpoint system only captures session context when agents push commits +to GitHub. This creates two gaps in session traceability: + +1. **Post-push context loss.** Agent reasoning, tool calls, and decisions made + after the final push are never captured. The checkpoint reflects the session + state at push time, not at session end. + +2. **No-push sessions have zero checkpoints.** Sessions that don't produce code + changes — review bots, research tasks, documentation-only work, sessions + that error before completion — leave no trace in the checkpoint system. + +The result: there is no way to query what an agent did during sessions that +didn't push, and even for sessions that did push, the final reasoning context +is lost. + +## Current Architecture + +### Checkpoint Models (`shared/egg_contracts/checkpoints.py`) + +The v1 schema defines: + +- **`Checkpoint`** (line 131): Full session context tied to a commit. `commit_sha` + is *required* — every checkpoint must reference a commit. Key fields: `id`, + `commit_sha`, `session` (SessionMetadata), `transcript`, `files_touched`, + `tool_calls`, `token_usage`, `issue_number`, `pr_number`, `pipeline_phase`, + `branch`, `push_sha`, `created_at`. + +- **`SessionMetadata`** (line 39): Session info including `session_id`, + `container_id`, `agent_role`, `started_at`. Also defines `ended_at` and + `duration_seconds` fields — but these are **never populated** in v1. + +- **`CheckpointSummary`** (line 215): Lightweight index entry. Also requires + `commit_sha`. + +- **`CheckpointIndex`** (line 252): Root-level index with flat list of + summaries and three lookup methods: `get_by_commit()`, `get_by_issue()`, + `get_by_branch()`. All lookups are O(n) linear scans. + +### Checkpoint Capture (`gateway/checkpoint_handler.py`) + +- **Single trigger point**: Checkpoints are only created on successful git push + (`capture_and_store_checkpoints_for_push`, line 747). + +- **Per-commit granularity**: One checkpoint per commit in multi-commit pushes, + all sharing the same transcript buffer content. + +- **Transcript source**: API proxy buffer at `/tmp/egg-transcripts/{container_id}.jsonl`, + written by `transcript_buffer.py` during Anthropic API proxying. + +- **Metadata resolution**: `issue_number`, `pr_number`, and `pipeline_phase` + are resolved from environment variables (`EGG_ISSUE_NUMBER`, `EGG_PR_NUMBER`, + `EGG_PIPELINE_PHASE`) at capture time (lines 330-349). + +- **Storage**: Checkpoints are stored on orphan branch `egg/checkpoints/v1` via + a temporary git worktree. Each checkpoint is JSON in a sharded directory + (`checkpoints/{2-char prefix}/ckpt-{id}.json`), plus an `index.json` at root. + Storage is async (background thread) to avoid blocking the push response. + +- **ID generation**: `generate_checkpoint_id_from_commit()` uses SHA-256 of + `{commit_sha}:{session_id}:{timestamp}` — deterministic but requires a + commit SHA. + +### Session Lifecycle (`gateway/session_manager.py`) + +The `Session` dataclass (line 94) tracks: `session_token`, `session_token_hash`, +`container_id`, `container_ip`, `mode`, `created_at`, `last_seen`, `expires_at`, +`agent_role`, `phase`. Notably, it does **not** store `issue_number` or +`pr_number`. + +Four deletion paths exist where session-end checkpoints should be captured: + +1. **`delete_session(token)`** (line 597): Deletes by token. Does *not* call + `_cleanup_transcript_buffer()` or capture any checkpoint. + +2. **`delete_session_by_container(container_id)`** (line 629): Deletes by + container ID. Calls `_cleanup_transcript_buffer()` which destroys the + transcript data. No checkpoint captured. + +3. **`prune_expired_sessions()`** (line 665): Removes all expired sessions + (24-hour TTL). Calls `_cleanup_transcript_buffer()` per session. No + checkpoint captured. + +4. **`cleanup_orphaned_worktrees(active_containers)`** in `worktree_manager.py` + (line 658): Removes worktrees for crashed containers. Receives only a set + of active container IDs — has no access to Session objects or transcript + buffers. + +### Transcript Size Limits (`shared/egg_contracts/transcript_extractor.py`) + +Size limits are defined as function parameter defaults: +- `max_content_length = 10,000` chars (message content) +- `max_param_length = 1,000` chars (tool parameters) +- `max_result_length = 500` chars (tool result summaries) +- `MAX_TRANSCRIPT_SIZE = 1,000,000` bytes (1MB, in `checkpoint_handler.py:104`) + +## Constraints + +**Schema constraints:** +- v2 schema must live on a separate branch (`egg/checkpoints/v2`) — no + migration needed, v1 remains untouched +- `commit_sha` must become optional for session-end checkpoints +- v2 models must coexist with v1 models in the same Python module + +**Session metadata gap:** +- `Session` does not store `issue_number` or `pr_number` +- Environment variables (`EGG_ISSUE_NUMBER`, etc.) are set in the agent + container, not the gateway process — they may not be accessible when + capturing session-end checkpoints from the gateway side +- Orphan cleanup has even less context: only container IDs, no Session objects + +**Transcript buffer lifecycle:** +- Buffer is destroyed by `_cleanup_transcript_buffer()` during session deletion +- Session-end checkpoint capture must read the buffer *before* cleanup +- For crashed containers, the buffer may be incomplete or already cleaned up + +**Concurrency:** +- Index file is updated per checkpoint write (atomic temp+rename pattern) +- With session-end checkpoints increasing volume, concurrent writes become + more likely but remain safe due to single-threaded index updates within + each `store_checkpoint` call + +**Performance:** +- Checkpoint storage involves git operations (worktree create, commit, push, + worktree remove) — adds latency to session deletion if synchronous +- Current push checkpoints use async storage (background thread) to avoid + blocking + +## Implementation Approaches + +### Session Metadata Propagation + +#### Option A: Store Metadata at Registration (Recommended) + +Add optional `issue_number` and `pr_number` fields to the `Session` dataclass. +Populate them during `register_session()`. At checkpoint capture time, read +from the Session object. Fall back to environment variables if Session lacks +the fields (backward compatibility with existing sessions). + +**Pros:** Reliable metadata at all capture points. Simple lookup. Works for +all deletion paths including orphan cleanup (if Session is retrieved before +deletion). + +**Cons:** Requires `register_session()` API change. Existing sessions lack +these fields until re-registered. + +#### Option B: Parse from Transcript Buffer + +Extract issue/PR numbers from transcript content at capture time. + +**Pros:** No Session model changes. + +**Cons:** Unreliable — depends on transcript format. Fails for crashed/empty +sessions. Performance overhead for large transcripts. + +#### Option C: Accept Partial Linkage + +Session-end checkpoints may lack issue/PR linkage. Rely on commit-linked +checkpoints for workflow correlation. + +**Pros:** Simplest implementation. + +**Cons:** Defeats the purpose of rich querying for no-push sessions. Cannot +query failed/research sessions by issue. + +### Checkpoint Timing + +#### Option A: Synchronous Capture + +Block session deletion until checkpoint is captured and stored. + +**Pros:** Guarantees capture. Simple control flow. + +**Cons:** Adds git operation latency (~5-15s) to every session deletion. +Storage failure blocks deletion. + +#### Option B: Async with Buffer Preservation (Recommended) + +Start async checkpoint capture. Defer transcript buffer cleanup until capture +signals completion or times out (30 seconds). + +**Pros:** Non-blocking for the API caller. Buffer guaranteed available during +capture. Graceful degradation on timeout. + +**Cons:** More complex lifecycle coordination. Need event/semaphore mechanism. + +#### Option C: Fire-and-Forget + +Start async capture, immediately proceed with cleanup. + +**Pros:** Simplest. No latency impact. + +**Cons:** Race condition — buffer may be deleted before capture reads it. +No guarantee of success. + +### Crash Checkpoint Completeness + +#### Option A: Status Field Only + +Mark crashed sessions with `session_status=FAILED`. + +**Pros:** Simple. Status conveys the key information. + +**Cons:** No indication of transcript completeness. + +#### Option B: Both Status and Truncation Reason (Recommended) + +Use `session_status=FAILED` and set `truncated_reason="container_crash"` on the +Transcript when applicable. + +**Pros:** Maximum clarity. Status indicates terminal state, truncation +indicates data quality. Supports analytics on crash rates vs data completeness. +Leverages existing `Transcript.truncated` and `truncation_reason` fields. + +**Cons:** Slightly redundant — but serves different query patterns. + +### Index Strategy + +#### Option A: Per-Checkpoint Index Update (Recommended) + +Keep the current approach: update the index atomically on every checkpoint write. +Add secondary index dictionaries for O(1) lookups. + +**Pros:** Strong consistency. Simple mental model. Current volume is low enough +that contention is not a concern. + +**Cons:** May need revisiting at high scale. + +#### Option B: Batched/Eventual Consistency + +Buffer index updates and flush periodically. + +**Pros:** Reduced write contention. + +**Cons:** Stale index between flushes. More complex. Premature optimization +for current scale. + +## Recommendation + +**Session metadata**: Option A — store at registration with environment +variable fallback. This provides reliable linkage for all capture paths. + +**Checkpoint timing**: Option B — async with 30-second buffer preservation +timeout. Balances reliability with performance. + +**Crash handling**: Option B — both status and truncation reason. Maximizes +clarity for downstream consumers. + +**Index strategy**: Option A — per-checkpoint updates with secondary indices. +Current scale does not warrant batching complexity. + +The overall approach is a breaking change (v2 branch, no dual-write). The v1 +branch remains readable but no new checkpoints are written to it. This matches +the owner's stated preference for switching straight to v2. + +## Files to Modify + +| File | Changes | +|------|---------| +| `shared/egg_contracts/checkpoints.py` | Add `TriggerType`, `SessionStatus`, `AgentType` enums. Add `CheckpointV2`, `CheckpointSummaryV2`, `CheckpointIndexV2` models with optional `commit_sha`, required `trigger_type`/`session_id`, secondary index dicts | +| `shared/egg_contracts/checkpoint_loader.py` | Add `generate_checkpoint_id_v2()` (no commit required), `save_checkpoint_v2()`, `load_checkpoint_v2()`, `load_checkpoint_index_v2()`, `add_checkpoint_to_index_v2()` with multi-index updates | +| `gateway/session_manager.py` | Add `issue_number`/`pr_number` to `Session` dataclass and persistence. Update `register_session()`. Refactor `_cleanup_transcript_buffer()` to await async capture. Hook `delete_session()`, `delete_session_by_container()`, `prune_expired_sessions()` to capture session-end checkpoints | +| `gateway/checkpoint_handler.py` | Add `capture_session_end_checkpoint()`. Add `store_checkpoint_v2()` for `egg/checkpoints/v2` branch. Update `capture_checkpoint()` to produce v2 format. Update `CHECKPOINT_BRANCH` constant. Increase `MAX_TRANSCRIPT_SIZE` to 3MB | +| `gateway/worktree_manager.py` | Add session lookup before `cleanup_orphaned_worktrees()`. Capture `FAILED` checkpoints for orphaned containers | +| `shared/egg_contracts/transcript_extractor.py` | Increase `max_content_length` to 25,000, `max_param_length` to 2,500, `max_result_length` to 1,500 | + +## Testing Strategy + +**Unit tests:** +- v2 model validation (CheckpointV2 with/without commit_sha, TriggerType/SessionStatus enums) +- v2 loader functions (save/load, index updates with secondary indices) +- Session metadata persistence with new fields, backward compatibility + +**Integration tests:** +- Session deletion triggers checkpoint capture with `COMPLETED` status +- TTL expiration triggers checkpoint with `EXPIRED` status +- Orphan cleanup triggers checkpoint with `FAILED` status and truncation reason +- Push triggers v2 checkpoint with `COMMIT` trigger type +- Multi-index queries return correct results across all dimensions +- Async capture with 30s timeout: buffer preserved until capture completes +- Buffer cleanup blocked until async capture finishes or times out + +**Existing test commands:** +- `PYTHONPATH=shared pytest shared/egg_contracts/tests/` +- `PYTHONPATH=gateway:shared pytest gateway/tests/` + +## Open Questions + +### 1. Orphan Session Lookup + +`cleanup_orphaned_worktrees()` receives only container IDs, not Session objects. +To capture checkpoints for crashed sessions, we need to look up Session metadata +by container_id. The `get_session_by_container()` method exists but returns +`None` for already-deleted sessions. The recommended approach is to capture +checkpoint *before* session deletion in each path, ensuring the Session object +is still available. For orphan cleanup specifically, the worktree manager should +be given a reference to the session manager to look up sessions before cleanup. + +### 2. Retention Policy + +No TTL for v2 checkpoints initially. Pruning can be added later if storage +grows. The git branch format makes it easy to rewrite history or drop old +checkpoints. + +### 3. Transcript Size for Session-End Checkpoints + +Session-end checkpoints may represent longer sessions than commit checkpoints. +The proposed increase from 1MB to 3MB for `MAX_TRANSCRIPT_SIZE` and 2.5x +increases to content/parameter/result limits should accommodate this. diff --git a/.egg-state/drafts/530-plan.md b/.egg-state/drafts/530-plan.md new file mode 100644 index 0000000000..88f31dd7b4 --- /dev/null +++ b/.egg-state/drafts/530-plan.md @@ -0,0 +1,447 @@ +# Plan: Checkpoint v2 with Rich Querying + +> Issue: #530 | Phase: plan + +## Summary + +This plan implements a v2 checkpoint system that captures every agent session regardless of push activity, enables rich querying across multiple dimensions (session, issue, PR, commit, agent type, phase), and provides clear semantics distinguishing commit-linked vs session-end checkpoints. The implementation follows the decisions from refinement: 30-second async buffer timeout, immediate switch to v2 (no dual-write), per-checkpoint index updates, no retention TTL, and increased transcript size limits. + +## Implementation Phases + +### Phase 1: Core Contracts & Models + +**Goal**: Define v2 data models with new enums, checkpoint structure, and multi-dimensional index. + +**Tasks**: +- [TASK-1-1] Add `TriggerType`, `SessionStatus`, and `AgentType` enums to checkpoints.py — Acceptance: Enums defined with all values from spec, unit tests pass +- [TASK-1-2] Add `CheckpointV2` model with optional `commit_sha`, required `trigger_type` and `session_id`, and new timestamp fields — Acceptance: Model validates correctly, supports both commit and session-end checkpoints +- [TASK-1-3] Add `CheckpointSummaryV2` model with all queryable fields — Acceptance: Summary captures all fields needed for index lookups +- [TASK-1-4] Add `CheckpointIndexV2` model with secondary indices (by_session, by_issue, by_pr, by_commit, by_agent_type, by_phase, by_trigger, by_status) — Acceptance: Index model supports O(1) lookups for all dimensions +- [TASK-1-5] Add unit tests for v2 models — Acceptance: Tests cover validation, serialization, and edge cases + +**Dependencies**: None + +**Exit criteria**: All v2 models defined and tested, `make test` passes for checkpoint tests + +### Phase 2: Checkpoint Loader v2 Functions + +**Goal**: Add v2 save/load functions with multi-index update logic. + +**Tasks**: +- [TASK-2-1] Add `generate_checkpoint_id_v2()` function for session-end checkpoints (not commit-based) — Acceptance: Generates unique IDs without requiring commit_sha +- [TASK-2-2] Add `save_checkpoint_v2()` with atomic write support — Acceptance: Writes to v2 directory structure, atomic file operations +- [TASK-2-3] Add `load_checkpoint_v2()` and `load_checkpoint_index_v2()` functions — Acceptance: Correctly loads v2 checkpoints and index +- [TASK-2-4] Add `add_checkpoint_to_index_v2()` with multi-index updates — Acceptance: Updates all secondary indices atomically (by_session, by_issue, by_pr, etc.) +- [TASK-2-5] Add v2-specific lookup helpers (get_by_session, get_by_trigger, get_by_status) — Acceptance: O(1) lookups using secondary indices +- [TASK-2-6] Add unit tests for v2 loader functions — Acceptance: Tests cover save/load, index updates, and concurrent access patterns + +**Dependencies**: Phase 1 + +**Exit criteria**: All v2 loader functions implemented and tested + +### Phase 3: Session Metadata Enhancement + +**Goal**: Store issue/PR numbers at session registration for reliable checkpoint linkage. + +**Tasks**: +- [TASK-3-1] Add `issue_number` and `pr_number` optional fields to `Session` dataclass — Acceptance: Fields persisted to disk, backward compatible with existing sessions +- [TASK-3-2] Update `register_session()` to accept optional `issue_number` and `pr_number` parameters — Acceptance: Parameters stored in Session, passed to gateway API +- [TASK-3-3] Update `Session.to_dict_for_persistence()` and `from_persistence()` for new fields — Acceptance: New fields round-trip correctly through persistence +- [TASK-3-4] Add unit tests for session metadata — Acceptance: Tests cover new fields, backward compatibility + +**Dependencies**: None (can be done in parallel with Phase 1-2) + +**Exit criteria**: Session model extended, persistence tested + +### Phase 4: Transcript Size Limit Increase + +**Goal**: Increase transcript extraction size limits for session-end checkpoints capturing longer sessions. + +**Tasks**: +- [TASK-4-1] Increase `max_content_length` from 10,000 to 25,000 characters — Acceptance: Longer messages preserved, existing tests pass +- [TASK-4-2] Increase `max_param_length` from 1,000 to 2,500 characters — Acceptance: Tool parameters less aggressively truncated +- [TASK-4-3] Increase `max_result_length` from 500 to 1,500 characters — Acceptance: Tool results more complete +- [TASK-4-4] Increase `MAX_TRANSCRIPT_SIZE` in checkpoint_handler.py from 1MB to 3MB — Acceptance: Larger transcripts accepted +- [TASK-4-5] Update constants documentation — Acceptance: Comments reflect new limits and rationale + +**Dependencies**: None (can be done in parallel with Phase 1-3) + +**Exit criteria**: Size limits increased, tests pass + +### Phase 5: Session-End Checkpoint Capture + +**Goal**: Implement checkpoint capture for session termination with async buffer preservation. + +**Tasks**: +- [TASK-5-1] Add `capture_session_end_checkpoint()` function to checkpoint_handler.py — Acceptance: Captures checkpoint with `trigger_type=SESSION_END`, uses v2 schema +- [TASK-5-2] Add `store_checkpoint_v2()` method to `CheckpointHandler` for v2 branch (`egg/checkpoints/v2`) — Acceptance: Stores to v2 branch with proper index updates +- [TASK-5-3] Implement async capture with buffer preservation (30-second timeout) — Acceptance: Buffer not deleted until capture completes or timeout, non-blocking for API +- [TASK-5-4] Add `truncated_reason` field handling for crashed sessions — Acceptance: Transcript marked with `truncated_reason="container_crash"` when applicable +- [TASK-5-5] Update `capture_checkpoint()` to produce v2 checkpoints with `trigger_type=COMMIT` — Acceptance: Existing push-triggered checkpoints use v2 format + +**Dependencies**: Phases 1, 2, 3 + +**Exit criteria**: Session-end checkpoint capture implemented with proper async/timeout behavior + +### Phase 6: Integration with Session Deletion Paths + +**Goal**: Hook checkpoint capture into all session termination paths. + +**Tasks**: +- [TASK-6-1] Update `delete_session()` to capture checkpoint with `status=COMPLETED` before buffer cleanup — Acceptance: Graceful deletions produce checkpoints +- [TASK-6-2] Update `delete_session_by_container()` to capture checkpoint with `status=COMPLETED` — Acceptance: Container-based deletions produce checkpoints +- [TASK-6-3] Update `prune_expired_sessions()` to capture checkpoint with `status=EXPIRED` for each expired session — Acceptance: TTL expirations produce checkpoints +- [TASK-6-4] Add session lookup before worktree cleanup in `cleanup_orphaned_worktrees()` — Acceptance: Session metadata available for orphan checkpoints +- [TASK-6-5] Update `cleanup_orphaned_worktrees()` to capture checkpoint with `status=FAILED` — Acceptance: Crashed containers produce checkpoints with `truncated_reason="container_crash"` +- [TASK-6-6] Refactor `_cleanup_transcript_buffer()` to await checkpoint capture completion — Acceptance: Buffer preserved until capture finishes or 30s timeout + +**Dependencies**: Phases 3, 5 + +**Exit criteria**: All deletion paths capture checkpoints, buffer lifecycle correct + +### Phase 7: Update Push-Triggered Checkpoints + +**Goal**: Migrate push-triggered checkpoints to v2 format and branch. + +**Tasks**: +- [TASK-7-1] Update `CHECKPOINT_BRANCH` constant to `egg/checkpoints/v2` — Acceptance: New checkpoints written to v2 branch +- [TASK-7-2] Update `capture_and_store_checkpoint()` to use v2 capture and storage — Acceptance: Push checkpoints use v2 models +- [TASK-7-3] Update `capture_and_store_checkpoints_for_push()` to produce v2 checkpoints — Acceptance: Multi-commit pushes produce v2 checkpoints +- [TASK-7-4] Remove v1 imports and references (clean break) — Acceptance: No v1 code paths remain active + +**Dependencies**: Phases 1, 2, 5 + +**Exit criteria**: All push-triggered checkpoints use v2 format on v2 branch + +### Phase 8: Integration Testing & Documentation + +**Goal**: Ensure end-to-end functionality and document the v2 system. + +**Tasks**: +- [TASK-8-1] Add integration tests for session-end checkpoint flow — Acceptance: Test covers graceful deletion, TTL expiration, and crash scenarios +- [TASK-8-2] Add integration tests for commit checkpoint flow with v2 — Acceptance: Push-triggered checkpoints work with v2 +- [TASK-8-3] Add integration tests for multi-dimensional index queries — Acceptance: Queries by session, issue, PR, status all work +- [TASK-8-4] Update checkpoint docstrings and module-level documentation — Acceptance: Code is well-documented for future maintainers +- [TASK-8-5] Verify no regressions in existing gateway functionality — Acceptance: Full test suite passes, manual smoke test succeeds + +**Dependencies**: All previous phases + +**Exit criteria**: All tests pass, documentation complete + +## Test Strategy + +- **Unit tests**: + - v2 model validation (CheckpointV2, CheckpointSummaryV2, CheckpointIndexV2) + - Enum values and serialization + - Loader functions (save/load, index updates) + - Session metadata persistence with new fields + +- **Integration tests**: + - Session deletion triggers checkpoint capture + - TTL expiration triggers checkpoint capture + - Orphan cleanup triggers checkpoint capture with FAILED status + - Push triggers v2 checkpoint capture + - Multi-index queries return correct results + - Async capture with timeout works correctly + - Buffer preservation during async capture + +- **Manual testing**: + - Start a container, make changes, gracefully delete → verify checkpoint + - Let a session expire via TTL → verify checkpoint with EXPIRED status + - Simulate container crash → verify checkpoint with FAILED status + - Push commits → verify v2 checkpoint created + - Query checkpoints by various dimensions (issue, session, status) + +## Rollback Plan + +1. **Immediate rollback**: Revert the commit on `egg/issue-530` branch +2. **Branch-level isolation**: v2 checkpoints are on a separate branch (`egg/checkpoints/v2`), v1 remains untouched +3. **Feature flag fallback**: If needed, add `CHECKPOINT_V2_ENABLED` env var to disable v2 capture and fall back to v1-only behavior (would require minimal additional code) +4. **Data recovery**: Since v1 branch is unchanged, all existing checkpoint data remains accessible + +If rollback is needed: +```bash +# Revert the feature branch +git checkout main +git branch -D egg/issue-530 # or revert specific commits + +# v1 checkpoints continue to work unchanged +# v2 branch can be deleted if desired (data loss for v2 checkpoints only) +git push origin --delete egg/checkpoints/v2 +``` + +## Risk Assessment + +| Risk | Likelihood | Impact | Mitigation | +|------|------------|--------|------------| +| Session metadata not available for orphan cleanup | Medium | Medium | Pre-populate metadata at registration, graceful degradation to null linkage | +| Buffer deleted before capture completes | Low | High | Implement async capture with explicit buffer preservation and 30s timeout | +| Index file corruption from concurrent writes | Low | Medium | Atomic write pattern already in place, single-threaded index updates | +| Increased checkpoint storage costs | Medium | Low | No TTL requested; can add pruning later if needed | +| Breaking change impacts existing tooling | Medium | Medium | Separate v2 branch, no v1 modifications, clear documentation | +| Gateway crash during async capture loses checkpoint | Low | Medium | Fire timeout, log warning; checkpoint loss is acceptable for rare edge case | + +## Migration Notes + +- **No database migrations**: Checkpoints are stored as JSON files in git +- **No config changes**: Uses existing environment variables, new optional params at registration +- **Breaking change**: v2 checkpoints written to `egg/checkpoints/v2` branch, not backward compatible with v1 readers +- **Parallel operation**: v1 checkpoints remain in place on `egg/checkpoints/v1` branch, but no new checkpoints written there +- **Session persistence**: New fields (`issue_number`, `pr_number`) are optional, existing sessions load correctly + +--- + +## Structured Task Appendix + +The following YAML block is machine-readable and will be extracted into the contract. +It must accurately reflect the tasks described above. The `pr:` section provides the +title and description that will be used when creating the pull request. + +```yaml +# yaml-tasks +pr: + title: "Checkpoint v2: capture all sessions with rich querying" + description: | + Implements a v2 checkpoint system that captures every agent session regardless + of push activity. Adds session-end checkpoints (COMPLETED, EXPIRED, FAILED) + with multi-dimensional index for querying by session, issue, PR, commit, + agent type, phase, and status. Resolves #530. +phases: + - id: 1 + name: Core Contracts & Models + goal: Define v2 data models with new enums, checkpoint structure, and multi-dimensional index + tasks: + - id: TASK-1-1 + description: Add TriggerType, SessionStatus, and AgentType enums to checkpoints.py + acceptance: Enums defined with all values from spec, unit tests pass + files: + - shared/egg_contracts/checkpoints.py + - id: TASK-1-2 + description: Add CheckpointV2 model with optional commit_sha, required trigger_type and session_id + acceptance: Model validates correctly, supports both commit and session-end checkpoints + files: + - shared/egg_contracts/checkpoints.py + - id: TASK-1-3 + description: Add CheckpointSummaryV2 model with all queryable fields + acceptance: Summary captures all fields needed for index lookups + files: + - shared/egg_contracts/checkpoints.py + - id: TASK-1-4 + description: Add CheckpointIndexV2 model with secondary indices + acceptance: Index model supports O(1) lookups for all dimensions + files: + - shared/egg_contracts/checkpoints.py + - id: TASK-1-5 + description: Add unit tests for v2 models + acceptance: Tests cover validation, serialization, and edge cases + files: + - shared/egg_contracts/tests/test_checkpoints.py + - id: 2 + name: Checkpoint Loader v2 Functions + goal: Add v2 save/load functions with multi-index update logic + tasks: + - id: TASK-2-1 + description: Add generate_checkpoint_id_v2() function for session-end checkpoints + acceptance: Generates unique IDs without requiring commit_sha + files: + - shared/egg_contracts/checkpoint_loader.py + - id: TASK-2-2 + description: Add save_checkpoint_v2() with atomic write support + acceptance: Writes to v2 directory structure, atomic file operations + files: + - shared/egg_contracts/checkpoint_loader.py + - id: TASK-2-3 + description: Add load_checkpoint_v2() and load_checkpoint_index_v2() functions + acceptance: Correctly loads v2 checkpoints and index + files: + - shared/egg_contracts/checkpoint_loader.py + - id: TASK-2-4 + description: Add add_checkpoint_to_index_v2() with multi-index updates + acceptance: Updates all secondary indices atomically + files: + - shared/egg_contracts/checkpoint_loader.py + - id: TASK-2-5 + description: Add v2-specific lookup helpers (get_by_session, get_by_trigger, get_by_status) + acceptance: O(1) lookups using secondary indices + files: + - shared/egg_contracts/checkpoint_loader.py + - id: TASK-2-6 + description: Add unit tests for v2 loader functions + acceptance: Tests cover save/load, index updates, and concurrent access patterns + files: + - shared/egg_contracts/tests/test_checkpoint_loader.py + - id: 3 + name: Session Metadata Enhancement + goal: Store issue/PR numbers at session registration for reliable checkpoint linkage + tasks: + - id: TASK-3-1 + description: Add issue_number and pr_number optional fields to Session dataclass + acceptance: Fields persisted to disk, backward compatible with existing sessions + files: + - gateway/session_manager.py + - id: TASK-3-2 + description: Update register_session() to accept optional issue_number and pr_number + acceptance: Parameters stored in Session, passed to gateway API + files: + - gateway/session_manager.py + - id: TASK-3-3 + description: Update Session persistence methods for new fields + acceptance: New fields round-trip correctly through persistence + files: + - gateway/session_manager.py + - id: TASK-3-4 + description: Add unit tests for session metadata + acceptance: Tests cover new fields, backward compatibility + files: + - gateway/tests/test_session_manager.py + - id: 4 + name: Transcript Size Limit Increase + goal: Increase transcript extraction size limits for session-end checkpoints + tasks: + - id: TASK-4-1 + description: Increase max_content_length from 10,000 to 25,000 characters + acceptance: Longer messages preserved, existing tests pass + files: + - shared/egg_contracts/transcript_extractor.py + - id: TASK-4-2 + description: Increase max_param_length from 1,000 to 2,500 characters + acceptance: Tool parameters less aggressively truncated + files: + - shared/egg_contracts/transcript_extractor.py + - id: TASK-4-3 + description: Increase max_result_length from 500 to 1,500 characters + acceptance: Tool results more complete + files: + - shared/egg_contracts/transcript_extractor.py + - id: TASK-4-4 + description: Increase MAX_TRANSCRIPT_SIZE from 1MB to 3MB + acceptance: Larger transcripts accepted + files: + - gateway/checkpoint_handler.py + - id: TASK-4-5 + description: Update constants documentation + acceptance: Comments reflect new limits and rationale + files: + - shared/egg_contracts/transcript_extractor.py + - gateway/checkpoint_handler.py + - id: 5 + name: Session-End Checkpoint Capture + goal: Implement checkpoint capture for session termination with async buffer preservation + tasks: + - id: TASK-5-1 + description: Add capture_session_end_checkpoint() function + acceptance: Captures checkpoint with trigger_type=SESSION_END, uses v2 schema + files: + - gateway/checkpoint_handler.py + - id: TASK-5-2 + description: Add store_checkpoint_v2() method for v2 branch + acceptance: Stores to egg/checkpoints/v2 branch with proper index updates + files: + - gateway/checkpoint_handler.py + - id: TASK-5-3 + description: Implement async capture with buffer preservation (30s timeout) + acceptance: Buffer not deleted until capture completes or timeout + files: + - gateway/checkpoint_handler.py + - id: TASK-5-4 + description: Add truncated_reason field handling for crashed sessions + acceptance: Transcript marked with truncated_reason when applicable + files: + - gateway/checkpoint_handler.py + - id: TASK-5-5 + description: Update capture_checkpoint() to produce v2 checkpoints + acceptance: Existing push-triggered checkpoints use v2 format + files: + - gateway/checkpoint_handler.py + - id: 6 + name: Integration with Session Deletion Paths + goal: Hook checkpoint capture into all session termination paths + tasks: + - id: TASK-6-1 + description: Update delete_session() to capture checkpoint with status=COMPLETED + acceptance: Graceful deletions produce checkpoints + files: + - gateway/session_manager.py + - id: TASK-6-2 + description: Update delete_session_by_container() to capture checkpoint + acceptance: Container-based deletions produce checkpoints + files: + - gateway/session_manager.py + - id: TASK-6-3 + description: Update prune_expired_sessions() to capture checkpoints with status=EXPIRED + acceptance: TTL expirations produce checkpoints + files: + - gateway/session_manager.py + - id: TASK-6-4 + description: Add session lookup before worktree cleanup + acceptance: Session metadata available for orphan checkpoints + files: + - gateway/worktree_manager.py + - id: TASK-6-5 + description: Update cleanup_orphaned_worktrees() to capture checkpoint with status=FAILED + acceptance: Crashed containers produce checkpoints + files: + - gateway/worktree_manager.py + - id: TASK-6-6 + description: Refactor _cleanup_transcript_buffer() to await checkpoint completion + acceptance: Buffer preserved until capture finishes or 30s timeout + files: + - gateway/session_manager.py + - id: 7 + name: Update Push-Triggered Checkpoints + goal: Migrate push-triggered checkpoints to v2 format and branch + tasks: + - id: TASK-7-1 + description: Update CHECKPOINT_BRANCH constant to egg/checkpoints/v2 + acceptance: New checkpoints written to v2 branch + files: + - gateway/checkpoint_handler.py + - id: TASK-7-2 + description: Update capture_and_store_checkpoint() to use v2 capture and storage + acceptance: Push checkpoints use v2 models + files: + - gateway/checkpoint_handler.py + - id: TASK-7-3 + description: Update capture_and_store_checkpoints_for_push() to produce v2 checkpoints + acceptance: Multi-commit pushes produce v2 checkpoints + files: + - gateway/checkpoint_handler.py + - id: TASK-7-4 + description: Remove v1 imports and references + acceptance: No v1 code paths remain active + files: + - gateway/checkpoint_handler.py + - id: 8 + name: Integration Testing & Documentation + goal: Ensure end-to-end functionality and document the v2 system + tasks: + - id: TASK-8-1 + description: Add integration tests for session-end checkpoint flow + acceptance: Test covers graceful deletion, TTL expiration, and crash scenarios + files: + - gateway/tests/test_checkpoint_handler.py + - id: TASK-8-2 + description: Add integration tests for commit checkpoint flow with v2 + acceptance: Push-triggered checkpoints work with v2 + files: + - gateway/tests/test_checkpoint_handler.py + - id: TASK-8-3 + description: Add integration tests for multi-dimensional index queries + acceptance: Queries by session, issue, PR, status all work + files: + - gateway/tests/test_checkpoint_handler.py + - id: TASK-8-4 + description: Update checkpoint docstrings and module-level documentation + acceptance: Code is well-documented for future maintainers + files: + - shared/egg_contracts/checkpoints.py + - gateway/checkpoint_handler.py + - id: TASK-8-5 + description: Verify no regressions in existing gateway functionality + acceptance: Full test suite passes, manual smoke test succeeds + files: [] +``` + +--- + +*Authored-by: egg* diff --git a/.egg-state/reviews/530-plan-agent-design-review.json b/.egg-state/reviews/530-plan-agent-design-review.json new file mode 100644 index 0000000000..bba6f5df6e --- /dev/null +++ b/.egg-state/reviews/530-plan-agent-design-review.json @@ -0,0 +1,7 @@ +{ + "reviewer": "agent-design", + "verdict": "approved", + "summary": "No agent-mode design anti-patterns found. The plan specifies what to build and where, without pre-fetching file contents into prompts, requiring structured output for human consumption, introducing post-processing pipelines, micromanaging implementation steps, or relying on prompt-level security. The YAML task appendix is explicitly machine-targeted for pipeline extraction. The 40-task breakdown is engineering specification appropriate for a multi-file feature, not rigid procedural instructions for the implementing agent.", + "feedback": "", + "timestamp": "2026-02-13T18:35:00Z" +} diff --git a/.egg-state/reviews/530-plan-unified-review.json b/.egg-state/reviews/530-plan-unified-review.json new file mode 100644 index 0000000000..f008771759 --- /dev/null +++ b/.egg-state/reviews/530-plan-unified-review.json @@ -0,0 +1,7 @@ +{ + "reviewer": "unified", + "verdict": "approved", + "summary": "The plan is comprehensive, well-structured, and faithfully implements all recommendations from the analysis. All 40 tasks across 8 phases are specific, actionable, and have clear acceptance criteria. Dependencies are correctly identified with parallelizable phases (1, 3, 4) noted. The test strategy covers unit, integration, and manual testing with appropriate edge cases (crashes, TTL expiration, concurrent access, buffer lifecycle). The risk assessment is thorough with a practical rollback plan leveraging the separate v2 branch. Minor nit: Phase 4 tasks are overly granular (individual constant changes could be one task), and TASK-8-4's acceptance criteria is subjective, but neither warrants revision.", + "feedback": "", + "timestamp": "2026-02-13T18:30:00Z" +} diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 0c822df7ab..f297d2b31c 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -88,7 +88,7 @@ Checkpoints capture agent session context as first-class versioned data in Git. - **Token usage**: Input/output tokens and estimated costs - **Session metadata**: Session ID, agent role, model, duration -Checkpoints are stored in the `egg/checkpoints/v1` branch and linked to commits via checkpoint IDs in contract tasks and audit entries. This provides full traceability from requirements to implementation. +Checkpoints are stored in the `egg/checkpoints/v2` branch and linked to commits via checkpoint IDs in contract tasks and audit entries. This provides full traceability from requirements to implementation. **Checkpoint CLI** diff --git a/gateway/checkpoint_handler.py b/gateway/checkpoint_handler.py index e2cd2ae42b..c56de4d203 100644 --- a/gateway/checkpoint_handler.py +++ b/gateway/checkpoint_handler.py @@ -1,10 +1,11 @@ """ Checkpoint handler for the gateway sidecar. -Captures agent session context on successful git push operations and stores -checkpoints in a dedicated branch (egg/checkpoints/v1). Each commit gets -exactly one checkpoint, enabling precise traceability between code changes -and agent sessions. +Captures agent session context on git push operations and session termination, +storing checkpoints in a dedicated branch (egg/checkpoints/v2). Supports two +trigger types: +- COMMIT: Created per-commit on git push (one checkpoint per commit) +- SESSION_END: Created when a session terminates (graceful, expired, or crashed) Architecture Overview: ┌──────────────────┐ ┌───────────────────────┐ @@ -20,32 +21,28 @@ │ (API capture) │ │ extractor │ │ (storage) │ └───────────────────┘ └───────────────────┘ └───────────────────┘ -Checkpoint Flow: +Checkpoint Flow (push): 1. gateway.py receives git push, forwards to GitHub 2. On success, calls capture_and_store_checkpoints_for_push() 3. Enumerates commits in push via get_commits_in_push() - 4. For each commit, captures checkpoint with transcript from proxy buffer - 5. Stores checkpoint in egg/checkpoints/v1 branch (async) + 4. For each commit, captures v2 checkpoint with trigger_type=COMMIT + 5. Stores checkpoint in egg/checkpoints/v2 branch (async) + +Checkpoint Flow (session end): + 1. Session deletion/expiry/crash triggers capture_session_end_checkpoint() + 2. Extracts transcript from proxy buffer (if available) + 3. Creates v2 checkpoint with trigger_type=SESSION_END + 4. Stores checkpoint asynchronously with buffer preservation Transcript Source: Transcripts are extracted from the API proxy buffer at - /tmp/egg-transcripts/{container_id}.jsonl. This file is written by - transcript_buffer.py when proxying Anthropic API calls. Using the - API proxy layer provides a stable format independent of Claude Code - internals. - - The buffer captures: - - API request/response pairs (messages, tools, content) - - Token usage per request - - Tool calls and their results - - Streaming and non-streaming responses + /tmp/egg-transcripts/{container_id}.jsonl. Integration points: - Called from gateway.py after successful push -- Uses SessionManager for session context -- Uses transcript_extractor for proxy buffer data extraction +- Called from session_manager.py on session deletion/expiry +- Called from worktree_manager.py on orphan cleanup - Uses checkpoint_loader for atomic writes to checkpoint branch -- Uses transcript_buffer for API traffic capture """ import os @@ -62,16 +59,20 @@ sys.path.insert(0, str(_shared_path)) from egg_contracts.checkpoint_loader import ( - add_checkpoint_to_index, + add_checkpoint_to_index_v2, generate_checkpoint_id_from_commit, + generate_checkpoint_id_v2, get_checkpoint_path, - save_checkpoint, + save_checkpoint_v2, ) from egg_contracts.checkpoints import ( - Checkpoint, + AgentType, + CheckpointV2, SessionMetadata, + SessionStatus, ToolCall, Transcript, + TriggerType, ) from egg_contracts.redactor import Redactor, RedactorConfig from egg_contracts.transcript_extractor import ( @@ -94,18 +95,40 @@ logger = get_logger("gateway.checkpoint-handler") -# Checkpoint branch name -CHECKPOINT_BRANCH = "egg/checkpoints/v1" +# Checkpoint branch name — v2 uses a separate branch from v1 +CHECKPOINT_BRANCH = "egg/checkpoints/v2" # Index file name INDEX_FILE = "index.json" # Maximum transcript size before truncation (bytes) -MAX_TRANSCRIPT_SIZE = 1_000_000 # 1MB +# Set to 3MB to capture more complete context from longer sessions, +# especially for session-end checkpoints that span entire agent sessions. +MAX_TRANSCRIPT_SIZE = 3_000_000 # 3MB + +# Timeout for async session-end checkpoint capture (seconds). +# Buffer cleanup is blocked until capture completes or this timeout elapses. +SESSION_END_CAPTURE_TIMEOUT = 30 # Feature flag for enabling/disabling checkpoints CHECKPOINT_ENABLED = os.environ.get("CHECKPOINT_ENABLED", "true").lower() == "true" +# Mapping from agent_role strings to AgentType enum values +_ROLE_TO_AGENT_TYPE = { + "coder": AgentType.CODER, + "tester": AgentType.TESTER, + "documenter": AgentType.DOCUMENTER, + "integrator": AgentType.INTEGRATOR, + "reviewer": AgentType.REVIEWER, +} + + +def _resolve_agent_type(agent_role: str | None) -> AgentType: + """Map an agent_role string to the AgentType enum.""" + if not agent_role: + return AgentType.UNKNOWN + return _ROLE_TO_AGENT_TYPE.get(agent_role.lower(), AgentType.UNKNOWN) + def get_commits_in_push( repo_path: str, @@ -139,13 +162,9 @@ def get_commits_in_push( # Handle new branch case - old_sha is null SHA (all zeros) null_sha = "0" * 40 if old_sha == null_sha or not old_sha: - # New branch - just return the tip commit - # Could also use rev-list from root, but that's expensive for large branches return [new_sha] try: - # Get commits between old and new (exclusive of old, inclusive of new) - # --reverse gives us chronological order (oldest first) result = subprocess.run( ["git", "rev-list", "--reverse", f"{old_sha}..{new_sha}"], cwd=repo_path, @@ -157,9 +176,8 @@ def get_commits_in_push( if result.returncode == 0 and result.stdout.strip(): commits = result.stdout.strip().split("\n") - return [c for c in commits if c] # Filter empty strings + return [c for c in commits if c] - # If rev-list fails or returns empty (force push, etc.), fall back to new_sha logger.debug( "git rev-list returned empty or failed, falling back to tip commit", old_sha=old_sha[:7], @@ -196,8 +214,9 @@ class CheckpointHandler: """ Handles checkpoint capture and storage for agent sessions. - This class is called after successful git push operations to capture - the session context and store it in the checkpoint branch. + Supports two checkpoint types: + - COMMIT: Created per-commit on git push + - SESSION_END: Created when a session terminates """ def __init__( @@ -205,13 +224,6 @@ def __init__( github_token: str | None = None, redactor_config: RedactorConfig | None = None, ): - """ - Initialize the checkpoint handler. - - Args: - github_token: GitHub token for pushing to checkpoint branch - redactor_config: Optional redactor configuration - """ self._github_token = github_token self._redactor = Redactor(redactor_config) @@ -224,15 +236,9 @@ def capture_checkpoint( issue_number: int | None = None, pipeline_phase: str | None = None, push_sha: str | None = None, - ) -> Checkpoint | None: + ) -> CheckpointV2 | None: """ - Capture a checkpoint for a commit. - - This method: - 1. Finds the current Claude Code session file - 2. Extracts transcript, tool calls, and token usage - 3. Applies redaction for sensitive data - 4. Creates and returns a Checkpoint object + Capture a v2 checkpoint for a commit (trigger_type=COMMIT). Args: repo_path: Path to the repository @@ -244,17 +250,13 @@ def capture_checkpoint( push_sha: Optional tip commit SHA of the push Returns: - Checkpoint object if successful, None if capture fails - - Raises: - CheckpointError: If checkpoint capture fails critically + CheckpointV2 object if successful, None if capture fails """ if not CHECKPOINT_ENABLED: logger.debug("Checkpoint capture disabled by CHECKPOINT_ENABLED=false") return None try: - # Get container ID from session for proxy buffer lookup container_id = session.container_id if session else None if not container_id: @@ -263,8 +265,8 @@ def capture_checkpoint( repo_path=repo_path, commit_sha=commit_sha[:7], ) - # Create a minimal checkpoint without transcript - return self._create_minimal_checkpoint( + return self._create_minimal_checkpoint_v2( + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, branch=branch, session=session, @@ -273,7 +275,6 @@ def capture_checkpoint( push_sha=push_sha, ) - # Get proxy buffer path for this container buffer_path = get_proxy_buffer_path(container_id) if not buffer_path.exists(): @@ -283,8 +284,8 @@ def capture_checkpoint( buffer_path=str(buffer_path), commit_sha=commit_sha[:7], ) - # Create a minimal checkpoint without transcript - return self._create_minimal_checkpoint( + return self._create_minimal_checkpoint_v2( + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, branch=branch, session=session, @@ -293,7 +294,6 @@ def capture_checkpoint( push_sha=push_sha, ) - # Extract transcript data from proxy buffer try: ( session_metadata, @@ -308,7 +308,8 @@ def capture_checkpoint( error=str(e), buffer_path=str(buffer_path), ) - return self._create_minimal_checkpoint( + return self._create_minimal_checkpoint_v2( + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, branch=branch, session=session, @@ -321,58 +322,47 @@ def capture_checkpoint( transcript = self._redact_transcript(transcript) tool_calls = self._redact_tool_calls(tool_calls) - # Merge session info if provided + # Merge session info if session: session_metadata.container_id = session.container_id session_metadata.agent_role = session.agent_role - # Get issue number from environment if not provided - if issue_number is None: - issue_str = os.environ.get("EGG_ISSUE_NUMBER") - if issue_str: - try: - issue_number = int(issue_str) - except ValueError: - pass - - # Get pipeline phase from environment if not provided - if pipeline_phase is None: - pipeline_phase = os.environ.get("EGG_PIPELINE_PHASE") - - # Get PR number from environment if available - pr_number: int | None = None - pr_str = os.environ.get("EGG_PR_NUMBER") - if pr_str: - try: - pr_number = int(pr_str) - except ValueError: - pass + # Resolve issue/PR/phase from session or env + issue_number = self._resolve_issue_number(issue_number, session) + pipeline_phase = self._resolve_pipeline_phase(pipeline_phase, session) + pr_number = self._resolve_pr_number(session) + agent_type = _resolve_agent_type(session.agent_role if session else None) - # Generate checkpoint ID checkpoint_id = generate_checkpoint_id_from_commit( commit_sha, session_metadata.session_id ) - # Create checkpoint - checkpoint = Checkpoint( + now = datetime.now(UTC) + checkpoint = CheckpointV2( id=checkpoint_id, + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, + push_sha=push_sha or commit_sha, + branch=branch, + session_id=session_metadata.session_id, + issue_number=issue_number, + pr_number=pr_number, + agent_type=agent_type, + pipeline_phase=pipeline_phase, session=session_metadata, transcript=transcript, files_touched=file_operations, tool_calls=tool_calls, token_usage=token_usage, - issue_number=issue_number, - pipeline_phase=pipeline_phase, - branch=branch, - created_at=datetime.now(UTC), - push_sha=push_sha or commit_sha, - pr_number=pr_number, + created_at=now, + session_started_at=session_metadata.started_at, + session_ended_at=session_metadata.ended_at, ) logger.info( "Checkpoint captured", checkpoint_id=checkpoint_id, + trigger_type="commit", commit_sha=commit_sha[:7], message_count=transcript.message_count, tool_call_count=len(tool_calls), @@ -387,40 +377,211 @@ def capture_checkpoint( error=str(e), commit_sha=commit_sha[:7], ) - # Don't raise - checkpoint failure should not block push return None - def _create_minimal_checkpoint( + def capture_session_end_checkpoint( self, - commit_sha: str, - branch: str, + session: Session, + session_status: SessionStatus, + repo_path: str | None = None, + ) -> CheckpointV2 | None: + """ + Capture a checkpoint when a session terminates. + + Creates a checkpoint with trigger_type=SESSION_END regardless of whether + the session pushed any commits. This ensures every session is captured. + + Args: + session: The Session object being terminated + session_status: Why the session ended (COMPLETED, EXPIRED, FAILED) + repo_path: Optional repo path (for finding repo context) + + Returns: + CheckpointV2 object if successful, None if capture fails + """ + if not CHECKPOINT_ENABLED: + return None + + container_id = session.container_id + + try: + now = datetime.now(UTC) + agent_type = _resolve_agent_type(session.agent_role) + session_id = container_id + + # Try to extract transcript from proxy buffer + transcript = None + tool_calls: list[ToolCall] = [] + file_operations = [] + token_usage = None + session_metadata = SessionMetadata( + session_id=session_id, + container_id=container_id, + agent_role=session.agent_role, + started_at=session.created_at, + ended_at=now, + duration_seconds=(now - session.created_at).total_seconds(), + ) + + try: + buffer_path = get_proxy_buffer_path(container_id) + if buffer_path.exists(): + ( + extracted_metadata, + transcript, + tool_calls, + file_operations, + token_usage, + ) = extract_transcript_from_proxy_buffer(buffer_path, container_id) + + # Merge extracted metadata with session info + extracted_metadata.container_id = container_id + extracted_metadata.agent_role = session.agent_role + session_metadata = extracted_metadata + + # Apply redaction + transcript = self._redact_transcript(transcript) + tool_calls = self._redact_tool_calls(tool_calls) + except (TranscriptExtractError, ValueError): + logger.debug( + "No transcript available for session-end checkpoint", + container_id=container_id, + ) + + # Mark transcript as truncated for crashed sessions + if session_status == SessionStatus.FAILED and transcript is not None: + transcript = transcript.model_copy( + update={ + "truncated": True, + "truncation_reason": "container_crash", + } + ) + + checkpoint_id = generate_checkpoint_id_v2(session_id, now) + + checkpoint = CheckpointV2( + id=checkpoint_id, + trigger_type=TriggerType.SESSION_END, + session_status=session_status, + session_id=session_id, + issue_number=session.issue_number, + pr_number=session.pr_number, + agent_type=agent_type, + pipeline_phase=session.phase, + session=session_metadata, + transcript=transcript, + files_touched=file_operations, + tool_calls=tool_calls, + token_usage=token_usage, + created_at=now, + session_started_at=session.created_at, + session_ended_at=now, + ) + + logger.info( + "Session-end checkpoint captured", + checkpoint_id=checkpoint_id, + container_id=container_id, + session_status=session_status.value, + agent_type=agent_type.value, + message_count=transcript.message_count if transcript else 0, + ) + + return checkpoint + + except Exception as e: + logger.error( + "Session-end checkpoint capture failed", + error=str(e), + container_id=container_id, + session_status=session_status.value, + ) + return None + + def _create_minimal_checkpoint_v2( + self, + trigger_type: TriggerType, + commit_sha: str | None = None, + branch: str | None = None, session: Session | None = None, issue_number: int | None = None, pipeline_phase: str | None = None, push_sha: str | None = None, - ) -> Checkpoint: - """Create a minimal checkpoint without transcript data.""" + session_status: SessionStatus | None = None, + ) -> CheckpointV2: + """Create a minimal v2 checkpoint without transcript data.""" session_id = session.container_id if session else "unknown" - checkpoint_id = generate_checkpoint_id_from_commit(commit_sha, session_id) + now = datetime.now(UTC) + + if trigger_type == TriggerType.COMMIT and commit_sha: + checkpoint_id = generate_checkpoint_id_from_commit(commit_sha, session_id) + else: + checkpoint_id = generate_checkpoint_id_v2(session_id, now) session_metadata = SessionMetadata( session_id=session_id, container_id=session.container_id if session else None, agent_role=session.agent_role if session else None, - started_at=datetime.now(UTC), + started_at=session.created_at if session else now, ) - return Checkpoint( + agent_type = _resolve_agent_type(session.agent_role if session else None) + + return CheckpointV2( id=checkpoint_id, + trigger_type=trigger_type, + session_status=session_status, commit_sha=commit_sha, - session=session_metadata, + push_sha=push_sha or commit_sha, + branch=branch, + session_id=session_id, issue_number=issue_number, + pr_number=session.pr_number if session else None, + agent_type=agent_type, pipeline_phase=pipeline_phase, - branch=branch, - created_at=datetime.now(UTC), - push_sha=push_sha or commit_sha, + session=session_metadata, + created_at=now, + session_started_at=session.created_at if session else now, ) + def _resolve_issue_number( + self, issue_number: int | None, session: Session | None + ) -> int | None: + """Resolve issue number from args, session, or environment.""" + if issue_number is not None: + return issue_number + if session and session.issue_number is not None: + return session.issue_number + issue_str = os.environ.get("EGG_ISSUE_NUMBER") + if issue_str: + try: + return int(issue_str) + except ValueError: + pass + return None + + def _resolve_pipeline_phase( + self, pipeline_phase: str | None, session: Session | None + ) -> str | None: + """Resolve pipeline phase from args, session, or environment.""" + if pipeline_phase is not None: + return pipeline_phase + if session and session.phase is not None: + return session.phase + return os.environ.get("EGG_PIPELINE_PHASE") + + def _resolve_pr_number(self, session: Session | None) -> int | None: + """Resolve PR number from session or environment.""" + if session and session.pr_number is not None: + return session.pr_number + pr_str = os.environ.get("EGG_PR_NUMBER") + if pr_str: + try: + return int(pr_str) + except ValueError: + pass + return None + def _redact_transcript(self, transcript: Transcript) -> Transcript: """Redact sensitive data from transcript messages.""" redacted_messages = [] @@ -437,15 +598,12 @@ def _redact_tool_calls(self, tool_calls: list[ToolCall]) -> list[ToolCall]: """Redact sensitive data from tool calls.""" redacted_calls = [] for tc in tool_calls: - # Redact parameters redacted_params = self._redactor.redact_dict(tc.parameters) - # Redact result summary redacted_result = None if tc.result_summary: redacted_result = self._redactor.redact_text(tc.result_summary) - # Special handling for Bash commands if tc.name == "Bash" and "command" in redacted_params: redacted_params["command"] = self._redactor.redact_command( str(redacted_params["command"]) @@ -462,23 +620,20 @@ def _redact_tool_calls(self, tool_calls: list[ToolCall]) -> list[ToolCall]: return redacted_calls - def store_checkpoint( + def store_checkpoint_v2( self, - checkpoint: Checkpoint, + checkpoint: CheckpointV2, repo_path: str, remote: str = "origin", ) -> bool: """ - Store a checkpoint in the checkpoint branch. + Store a v2 checkpoint in the checkpoint branch. - This method: - 1. Creates the checkpoint branch if it doesn't exist - 2. Writes the checkpoint JSON to the branch - 3. Updates the index - 4. Pushes to the remote + Writes the checkpoint JSON and updates the v2 multi-dimensional index + on the egg/checkpoints/v2 branch. Args: - checkpoint: The checkpoint to store + checkpoint: The v2 checkpoint to store repo_path: Path to the repository remote: Git remote name @@ -489,20 +644,16 @@ def store_checkpoint( return False try: - # Create a temporary directory for the checkout with tempfile.TemporaryDirectory(prefix="checkpoint_") as temp_dir: temp_path = Path(temp_dir) - # Check if checkpoint branch exists branch_exists = self._branch_exists(repo_path, remote, CHECKPOINT_BRANCH) if branch_exists: - # Fetch the branch self._run_git( repo_path, ["fetch", remote, f"{CHECKPOINT_BRANCH}:{CHECKPOINT_BRANCH}"], ) - # Checkout the branch to temp directory self._run_git( repo_path, [ @@ -514,47 +665,50 @@ def store_checkpoint( ], ) else: - # Create orphan branch self._run_git( repo_path, ["worktree", "add", "--detach", str(temp_path)], ) - # Initialize as orphan self._run_git( str(temp_path), ["checkout", "--orphan", CHECKPOINT_BRANCH], ) - # Remove any files self._run_git(str(temp_path), ["rm", "-rf", "."], check=False) try: - # Get checkpoint path checkpoint_path = get_checkpoint_path(temp_path / "checkpoints", checkpoint.id) - # Save checkpoint - save_checkpoint(checkpoint, checkpoint_path) + save_checkpoint_v2(checkpoint, checkpoint_path) - # Update the index for fast lookups by commit SHA index_path = temp_path / INDEX_FILE - add_checkpoint_to_index(checkpoint, index_path) + add_checkpoint_to_index_v2(checkpoint, index_path) - # Stage both the checkpoint and index files self._run_git( str(temp_path), ["add", str(checkpoint_path.relative_to(temp_path)), INDEX_FILE], ) - # Commit (--no-verify skips pre-commit hooks which are - # irrelevant for the checkpoint branch) - commit_msg = ( - f"Add checkpoint {checkpoint.id} for commit {checkpoint.commit_sha[:7]}" - ) + # Build a descriptive commit message + if checkpoint.trigger_type == TriggerType.COMMIT: + commit_desc = ( + f"commit {checkpoint.commit_sha[:7]}" + if checkpoint.commit_sha + else "unknown commit" + ) + else: + status = ( + checkpoint.session_status.value + if checkpoint.session_status + else "unknown" + ) + commit_desc = f"session-end ({status})" + + commit_msg = f"Add checkpoint {checkpoint.id} for {commit_desc}" self._run_git( str(temp_path), ["commit", "--no-verify", "-m", commit_msg], ) - # Push (use longer timeout for large transcripts) self._run_git( str(temp_path), ["push", remote, f"HEAD:{CHECKPOINT_BRANCH}"], @@ -565,6 +719,7 @@ def store_checkpoint( "Checkpoint stored", checkpoint_id=checkpoint.id, branch=CHECKPOINT_BRANCH, + trigger_type=checkpoint.trigger_type.value, ) # Update usage aggregates (graceful degradation on failure) @@ -576,7 +731,6 @@ def store_checkpoint( checkpoint_id=checkpoint.id, ) except (UsageLoadError, UsageSaveError) as e: - # Log but don't fail - usage update is non-critical logger.warning( "Failed to update usage aggregates", error=str(e), @@ -586,7 +740,6 @@ def store_checkpoint( return True finally: - # Clean up worktree self._run_git( repo_path, ["worktree", "remove", "--force", str(temp_path)], @@ -601,6 +754,9 @@ def store_checkpoint( ) return False + # Keep store_checkpoint as alias for backward compatibility + store_checkpoint = store_checkpoint_v2 + def _branch_exists(self, repo_path: str, remote: str, branch: str) -> bool: """Check if a branch exists on the remote.""" try: @@ -626,16 +782,14 @@ def _run_git( """Run a git command.""" env = os.environ.copy() - # Add GitHub token if available if self._github_token: env["GIT_ASKPASS"] = "echo" env["GIT_USERNAME"] = "x-access-token" env["GIT_PASSWORD"] = self._github_token # SECURITY: Disable all git hooks. The checkpoint handler runs git commands - # internally for bookkeeping (storing checkpoints to the egg/checkpoints/v1 branch). + # internally for bookkeeping (storing checkpoints to the checkpoint branch). # Hooks from user repos must not execute in the gateway's trusted environment. - # See issue #58 for context on hook-based attacks. cmd = ["git", "-c", "core.hooksPath=/dev/null"] + args result = subprocess.run( @@ -679,33 +833,18 @@ def capture_and_store_checkpoint( push_sha: str | None = None, github_token: str | None = None, async_store: bool = True, -) -> Checkpoint | None: +) -> CheckpointV2 | None: """ - Convenience function to capture and store a checkpoint. + Capture and store a v2 checkpoint for a commit. This is the main entry point for checkpoint capture from gateway.py. By default, storage is done asynchronously to not block the push response. - - Args: - repo_path: Path to the repository - commit_sha: The commit SHA to checkpoint - branch: The branch where the commit was made - session: Optional Session object with container/role info - issue_number: Optional GitHub issue number - pipeline_phase: Optional SDLC pipeline phase - push_sha: Optional tip commit SHA of the push - github_token: Optional GitHub token for pushing - async_store: If True, store checkpoint asynchronously - - Returns: - Checkpoint object if capture successful, None otherwise """ if not CHECKPOINT_ENABLED: return None handler = get_checkpoint_handler(github_token) - # Capture checkpoint checkpoint = handler.capture_checkpoint( repo_path=repo_path, commit_sha=commit_sha, @@ -719,12 +858,11 @@ def capture_and_store_checkpoint( if not checkpoint: return None - # Store checkpoint (async or sync) if async_store: - # Store in background thread to not block push response + def _store_with_error_handling() -> None: try: - handler.store_checkpoint(checkpoint, repo_path) + handler.store_checkpoint_v2(checkpoint, repo_path) except Exception as e: logger.error( "Async checkpoint storage failed", @@ -739,7 +877,7 @@ def _store_with_error_handling() -> None: ) thread.start() else: - handler.store_checkpoint(checkpoint, repo_path) + handler.store_checkpoint_v2(checkpoint, repo_path) return checkpoint @@ -754,40 +892,16 @@ def capture_and_store_checkpoints_for_push( pipeline_phase: str | None = None, github_token: str | None = None, async_store: bool = True, -) -> list[Checkpoint]: +) -> list[CheckpointV2]: """ - Capture and store checkpoints for all commits in a push. - - This function creates one checkpoint per commit in the push, enabling - per-commit granularity for traceability. The push_sha field on each - checkpoint points to the tip commit of the push (new_sha). - - Args: - repo_path: Path to the repository - old_sha: SHA before the push (remote ref before update) - new_sha: SHA after the push (new tip of the branch) - branch: The branch where the commit was made - session: Optional Session object with container/role info - issue_number: Optional GitHub issue number - pipeline_phase: Optional SDLC pipeline phase - github_token: Optional GitHub token for pushing - async_store: If True, store checkpoints asynchronously - - Returns: - List of Checkpoint objects that were successfully captured. - May be empty if checkpoints are disabled or capture fails. + Capture and store v2 checkpoints for all commits in a push. - Note: - All checkpoints in a multi-commit push share the same transcript buffer, - so each checkpoint contains the full session transcript up to that point. - This is intentional: the transcript captures the entire agent reasoning - session, not just the work for a specific commit. Deduplication of transcript - content is left to the storage layer if needed. + Creates one checkpoint per commit with trigger_type=COMMIT. The push_sha + field on each checkpoint points to the tip commit (new_sha). """ if not CHECKPOINT_ENABLED: return [] - # Get all commits in the push (oldest first) commits = get_commits_in_push(repo_path, old_sha, new_sha) if not commits: @@ -801,7 +915,7 @@ def capture_and_store_checkpoints_for_push( branch=branch, ) - checkpoints = [] + checkpoints: list[CheckpointV2] = [] handler = get_checkpoint_handler(github_token) for commit_sha in commits: @@ -813,7 +927,7 @@ def capture_and_store_checkpoints_for_push( session=session, issue_number=issue_number, pipeline_phase=pipeline_phase, - push_sha=new_sha, # All checkpoints point to the push tip + push_sha=new_sha, ) if checkpoint: @@ -824,25 +938,23 @@ def capture_and_store_checkpoints_for_push( commit_sha=commit_sha[:7], error=str(e), ) - # Continue with other commits - don't fail the whole push if not checkpoints: logger.debug("No checkpoints captured for push", push_sha=new_sha[:7]) return [] - # Store all checkpoints (async or sync) if async_store: - # Store in background thread to not block push response + def _store_all_with_error_handling() -> None: - for checkpoint in checkpoints: + for cp in checkpoints: try: - handler.store_checkpoint(checkpoint, repo_path) + handler.store_checkpoint_v2(cp, repo_path) except Exception as e: logger.error( "Async checkpoint storage failed", error=str(e), - checkpoint_id=checkpoint.id, - commit_sha=checkpoint.commit_sha[:7], + checkpoint_id=cp.id, + commit_sha=cp.commit_sha[:7] if cp.commit_sha else "none", ) thread = threading.Thread( @@ -851,7 +963,90 @@ def _store_all_with_error_handling() -> None: ) thread.start() else: - for checkpoint in checkpoints: - handler.store_checkpoint(checkpoint, repo_path) + for cp in checkpoints: + handler.store_checkpoint_v2(cp, repo_path) return checkpoints + + +def capture_session_end_checkpoint( + session: Session, + session_status: SessionStatus, + repo_path: str | None = None, + github_token: str | None = None, + async_store: bool = True, +) -> tuple[CheckpointV2 | None, threading.Event | None]: + """ + Capture and store a session-end checkpoint. + + Returns a (checkpoint, completion_event) tuple. The completion_event + is set when async storage finishes (or immediately if async_store=False). + Callers should wait on the event before cleaning up the transcript buffer. + + Args: + session: The Session being terminated + session_status: Terminal status (COMPLETED, EXPIRED, FAILED) + repo_path: Optional repo path for context + github_token: Optional GitHub token for pushing + async_store: If True, store asynchronously (default) + + Returns: + Tuple of (checkpoint, completion_event). completion_event is None + if async_store is False or checkpoint capture failed. + """ + if not CHECKPOINT_ENABLED: + return None, None + + handler = get_checkpoint_handler(github_token) + + checkpoint = handler.capture_session_end_checkpoint( + session=session, + session_status=session_status, + repo_path=repo_path, + ) + + if not checkpoint: + return None, None + + # Determine repo_path for storage + if repo_path is None: + # Try to find a repo path from known locations + repos_base = Path("/home/egg/repos") + if repos_base.exists(): + for repo_dir in repos_base.iterdir(): + if repo_dir.is_dir() and (repo_dir / ".git").exists(): + repo_path = str(repo_dir) + break + + if repo_path is None: + logger.warning( + "No repo path available for session-end checkpoint storage", + container_id=session.container_id, + ) + return checkpoint, None + + if async_store: + completion_event = threading.Event() + + def _store_with_event() -> None: + try: + handler.store_checkpoint_v2(checkpoint, repo_path) + except Exception as e: + logger.error( + "Async session-end checkpoint storage failed", + error=str(e), + checkpoint_id=checkpoint.id, + container_id=session.container_id, + ) + finally: + completion_event.set() + + thread = threading.Thread( + target=_store_with_event, + daemon=True, + ) + thread.start() + return checkpoint, completion_event + else: + handler.store_checkpoint_v2(checkpoint, repo_path) + return checkpoint, None diff --git a/gateway/gateway.py b/gateway/gateway.py index 82e87c58cf..65c651499a 100644 --- a/gateway/gateway.py +++ b/gateway/gateway.py @@ -563,7 +563,7 @@ def git_push() -> tuple[Response, int] | Response: # Checkpoint branch bypass: pushes to the checkpoint branch always succeed # regardless of session mode or phase (checkpoints can be created at any time) - CHECKPOINT_BRANCH = "egg/checkpoints/v1" + CHECKPOINT_BRANCH = "egg/checkpoints/v2" is_checkpoint_push = branch == CHECKPOINT_BRANCH repo_info = parse_owner_repo(repo) @@ -3730,7 +3730,10 @@ def main() -> None: # Clean up orphaned worktrees from crashed containers try: - orphans_removed = startup_cleanup(active_containers=active_container_ids) + orphans_removed = startup_cleanup( + active_containers=active_container_ids, + session_manager=get_session_manager(), + ) if orphans_removed > 0: logger.info(f"Startup cleanup removed {orphans_removed} orphaned worktree(s)") except Exception as e: diff --git a/gateway/phase_filter.py b/gateway/phase_filter.py index 621c42a131..e9ec5531ff 100644 --- a/gateway/phase_filter.py +++ b/gateway/phase_filter.py @@ -628,7 +628,9 @@ def check_phase_file_restrictions( restrictions = self._phase_file_restrictions.get(phase) if not restrictions: - return FileRestrictionResult.allow(f"No phase file restrictions for phase: {phase.value}") + return FileRestrictionResult.allow( + f"No phase file restrictions for phase: {phase.value}" + ) blocked_files: list[str] = [] blocked_reasons: list[str] = [] diff --git a/gateway/policy.py b/gateway/policy.py index c1a1a079c4..d731ecbfb2 100644 --- a/gateway/policy.py +++ b/gateway/policy.py @@ -882,7 +882,9 @@ def check_merge_allowed(self, repo: str, pr_number: int) -> PolicyResult: }, ) - def check_pr_review_allowed(self, repo: str, pr_number: int, auth_mode: str = "bot") -> PolicyResult: + def check_pr_review_allowed( + self, repo: str, pr_number: int, auth_mode: str = "bot" + ) -> PolicyResult: """ Check if posting a PR review is allowed. diff --git a/gateway/session_manager.py b/gateway/session_manager.py index 073e16816b..85f4eafef4 100644 --- a/gateway/session_manager.py +++ b/gateway/session_manager.py @@ -13,6 +13,8 @@ - Rate limiting prevents enumeration attacks """ +from __future__ import annotations + import hashlib import json import os @@ -52,6 +54,76 @@ def _cleanup_transcript_buffer(container_id: str) -> None: ) +_captured_containers: set[str] = set() +_captured_containers_lock = threading.Lock() + + +def _capture_and_cleanup_session( + session: Session, + session_status: str, +) -> None: + """Capture session-end checkpoint, then clean up transcript buffer. + + Ensures the transcript buffer is preserved until checkpoint capture completes. + Falls back to immediate cleanup if checkpoint capture is unavailable. + + Uses a per-container deduplication guard to prevent multiple captures from + racing code paths (delete_session, cleanup_orphaned_worktrees, prune_expired). + """ + # Deduplicate: only capture once per container + with _captured_containers_lock: + if session.container_id in _captured_containers: + logger.debug( + "Session-end checkpoint already captured, skipping", + container_id=session.container_id, + session_status=session_status, + ) + return + _captured_containers.add(session.container_id) + + try: + from checkpoint_handler import ( + SESSION_END_CAPTURE_TIMEOUT, + capture_session_end_checkpoint, + ) + from egg_contracts.checkpoints import SessionStatus + + status = ( + session_status + if isinstance(session_status, SessionStatus) + else SessionStatus(session_status) + ) + _checkpoint, completion_event = capture_session_end_checkpoint( + session=session, + session_status=status, + ) + + # Wait for async storage to complete before cleaning up the buffer + if completion_event is not None: + completion_event.wait(timeout=SESSION_END_CAPTURE_TIMEOUT) + + except ImportError: + logger.debug( + "checkpoint_handler not available, skipping session-end checkpoint", + container_id=session.container_id, + ) + except Exception as e: + logger.warning( + "Session-end checkpoint capture failed", + container_id=session.container_id, + error=str(e), + ) + finally: + # Always clean up the buffer after checkpoint capture + _cleanup_transcript_buffer(session.container_id) + + # Remove from dedup set to prevent unbounded growth. + # The session is fully processed (checkpoint captured, buffer cleaned up) + # and no future code path will call this for the same container. + with _captured_containers_lock: + _captured_containers.discard(session.container_id) + + # Session configuration DEFAULT_SESSION_TTL_HOURS = 24 DEFAULT_CLEANUP_INTERVAL_MINUTES = 15 @@ -118,6 +190,8 @@ class Session: expires_at: datetime agent_role: str | None = None # Role set by workflow context phase: str | None = None # SDLC pipeline phase for operation filtering + issue_number: int | None = None # GitHub issue number for checkpoint linkage + pr_number: int | None = None # GitHub PR number for checkpoint linkage def is_expired(self) -> bool: """Check if session has expired.""" @@ -143,10 +217,14 @@ def to_dict_for_persistence(self) -> dict[str, Any]: result["agent_role"] = self.agent_role if self.phase is not None: result["phase"] = self.phase + if self.issue_number is not None: + result["issue_number"] = self.issue_number + if self.pr_number is not None: + result["pr_number"] = self.pr_number return result @classmethod - def from_persistence(cls, data: dict[str, Any]) -> "Session": + def from_persistence(cls, data: dict[str, Any]) -> Session: """Create Session from persisted data (no raw token).""" return cls( session_token=None, # Raw token not persisted @@ -159,6 +237,8 @@ def from_persistence(cls, data: dict[str, Any]) -> "Session": expires_at=datetime.fromisoformat(data["expires_at"]), agent_role=data.get("agent_role"), phase=data.get("phase"), + issue_number=data.get("issue_number"), + pr_number=data.get("pr_number"), ) @@ -301,6 +381,8 @@ def register_session( container_ip: str, mode: ModeType, phase: str | None = None, + issue_number: int | None = None, + pr_number: int | None = None, ) -> tuple[str, Session]: """ Register a new session for a container. @@ -310,6 +392,8 @@ def register_session( container_ip: Container's IP address on the Docker network mode: Repository visibility mode (private or public) phase: SDLC pipeline phase (e.g., "refine", "plan", "implement", "pr") + issue_number: Optional GitHub issue number for checkpoint linkage + pr_number: Optional GitHub PR number for checkpoint linkage Returns: Tuple of (session_token, Session) @@ -329,6 +413,8 @@ def register_session( last_seen=now, expires_at=now + timedelta(hours=self._ttl_hours), phase=phase, + issue_number=issue_number, + pr_number=pr_number, ) with self._lock: @@ -598,6 +684,7 @@ def delete_session(self, token: str) -> bool: """ Delete a session by token. + Captures a session-end checkpoint (COMPLETED) before cleaning up. Only the launcher (with launcher_secret) should call this. Args: @@ -607,6 +694,7 @@ def delete_session(self, token: str) -> bool: True if session was deleted, False if not found """ token_hash = self._token_to_hash.get(token) or _hash_token(token) + session = None with self._lock: session = self._sessions.get(token_hash) @@ -617,61 +705,76 @@ def delete_session(self, token: str) -> bool: self._token_to_hash.pop(token, None) self._save_to_disk() - logger.info( - "Session deleted", - event_type="session_deleted", - session_token_hash=token_hash[:16], - container_id=session.container_id, - ) + # Capture session-end checkpoint outside the lock to avoid + # blocking other session operations during the up-to-30s wait + _capture_and_cleanup_session(session, "completed") - return True + logger.info( + "Session deleted", + event_type="session_deleted", + session_token_hash=token_hash[:16], + container_id=session.container_id, + ) + + return True def delete_session_by_container(self, container_id: str) -> bool: """ Delete session by container ID. + Captures a session-end checkpoint (COMPLETED) before cleaning up. + Args: container_id: Docker container ID Returns: True if session was deleted, False if not found """ + session = None + token_hash = None + with self._lock: - to_delete = None - for token_hash, session in self._sessions.items(): - if session.container_id == container_id: - to_delete = token_hash + for th, s in self._sessions.items(): + if s.container_id == container_id: + token_hash = th + session = s break - if to_delete: - session = self._sessions.pop(to_delete) + if token_hash: + self._sessions.pop(token_hash) if session.session_token: self._token_to_hash.pop(session.session_token, None) self._save_to_disk() - # Clean up transcript buffer for this container - _cleanup_transcript_buffer(container_id) + if session: + # Capture session-end checkpoint outside the lock to avoid + # blocking other session operations during the up-to-30s wait + _capture_and_cleanup_session(session, "completed") - logger.info( - "Session deleted by container ID", - event_type="session_deleted", - session_token_hash=to_delete[:16], - container_id=container_id, - ) - return True + logger.info( + "Session deleted by container ID", + event_type="session_deleted", + session_token_hash=token_hash[:16], + container_id=container_id, + ) + return True - return False + return False def prune_expired_sessions(self) -> int: """ Remove all expired sessions. + Captures session-end checkpoints (EXPIRED) for each pruned session + before cleaning up transcript buffers. + Called periodically and on gateway startup. Returns: Number of sessions pruned """ - pruned = 0 + expired_sessions: list[tuple[str, Session]] = [] + with self._lock: expired_hashes = [ token_hash for token_hash, session in self._sessions.items() if session.is_expired() @@ -681,22 +784,43 @@ def prune_expired_sessions(self) -> int: session = self._sessions.pop(token_hash) if session.session_token: self._token_to_hash.pop(session.session_token, None) - pruned += 1 + expired_sessions.append((token_hash, session)) - # Clean up transcript buffer for expired session - _cleanup_transcript_buffer(session.container_id) + if expired_sessions: + self._save_to_disk() - logger.info( - "Session expired and pruned", - event_type="session_expired", - session_token_hash=token_hash[:16], - container_id=session.container_id, + # Capture checkpoints concurrently to avoid blocking N×30s sequentially. + # Each capture waits up to 30s for async storage, so we use threads. + if expired_sessions: + threads = [] + for token_hash, session in expired_sessions: + t = threading.Thread( + target=_capture_and_cleanup_session, + args=(session, "expired"), + daemon=True, ) + t.start() + threads.append((t, token_hash, session)) - if pruned > 0: - self._save_to_disk() + for t, token_hash, session in threads: + t.join(timeout=35) # 30s capture timeout + 5s buffer + if t.is_alive(): + logger.warning( + "Session expired, checkpoint capture still in progress", + event_type="session_expired", + session_token_hash=token_hash[:16], + container_id=session.container_id, + capture_timed_out=True, + ) + else: + logger.info( + "Session expired and pruned", + event_type="session_expired", + session_token_hash=token_hash[:16], + container_id=session.container_id, + ) - return pruned + return len(expired_sessions) def list_sessions(self) -> list[dict[str, Any]]: """ diff --git a/gateway/tests/test_checkpoint_handler.py b/gateway/tests/test_checkpoint_handler.py index c34f209491..a8cc19e6f6 100644 --- a/gateway/tests/test_checkpoint_handler.py +++ b/gateway/tests/test_checkpoint_handler.py @@ -1,17 +1,15 @@ -"""Tests for checkpoint_handler module - per-commit checkpoint creation.""" +"""Tests for checkpoint_handler module - per-commit and session-end checkpoint creation.""" import subprocess -from datetime import UTC, datetime -from unittest.mock import MagicMock, patch, call - -import pytest +from datetime import UTC, datetime, timedelta +from unittest.mock import MagicMock, patch # Import from conftest-loaded modules from checkpoint_handler import ( + capture_session_end_checkpoint, get_commits_in_push, - capture_and_store_checkpoints_for_push, - CheckpointHandler, ) +from session_manager import Session, _hash_token class TestGetCommitsInPush: @@ -177,7 +175,7 @@ def test_returns_empty_for_no_commits(self, mock_get_commits): @patch("checkpoint_handler.get_checkpoint_handler") def test_creates_checkpoint_for_each_commit(self, mock_get_handler, mock_get_commits): """Test that a checkpoint is created for each commit in the push.""" - from egg_contracts.checkpoints import Checkpoint, SessionMetadata + from egg_contracts.checkpoints import CheckpointV2, SessionMetadata, TriggerType # Mock 3 commits commit1 = "1111111111111111111111111111111111111111" @@ -191,19 +189,29 @@ def test_creates_checkpoint_for_each_commit(self, mock_get_handler, mock_get_com # Create mock checkpoints def create_checkpoint( - repo_path, commit_sha, branch, session=None, issue_number=None, pipeline_phase=None, push_sha=None + repo_path, + commit_sha, + branch, + session=None, + issue_number=None, + pipeline_phase=None, + push_sha=None, ): - return Checkpoint( + now = datetime.now(UTC) + return CheckpointV2( id=f"ckpt-{commit_sha[:12]}", + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, - session=SessionMetadata(session_id="test", started_at=datetime.now(UTC)), - created_at=datetime.now(UTC), + session_id="test", + session=SessionMetadata(session_id="test", started_at=now), + created_at=now, + session_started_at=now, push_sha=push_sha, branch=branch, ) mock_handler.capture_checkpoint.side_effect = create_checkpoint - mock_handler.store_checkpoint.return_value = True + mock_handler.store_checkpoint_v2.return_value = True import checkpoint_handler @@ -218,13 +226,13 @@ def create_checkpoint( # Should have 3 checkpoints assert len(checkpoints) == 3 assert mock_handler.capture_checkpoint.call_count == 3 - assert mock_handler.store_checkpoint.call_count == 3 + assert mock_handler.store_checkpoint_v2.call_count == 3 @patch("checkpoint_handler.get_commits_in_push") @patch("checkpoint_handler.get_checkpoint_handler") def test_push_sha_set_on_all_checkpoints(self, mock_get_handler, mock_get_commits): """Test that push_sha is set to the tip commit for all checkpoints.""" - from egg_contracts.checkpoints import Checkpoint, SessionMetadata + from egg_contracts.checkpoints import CheckpointV2, SessionMetadata, TriggerType # Mock 2 commits commit1 = "1111111111111111111111111111111111111111" @@ -235,26 +243,36 @@ def test_push_sha_set_on_all_checkpoints(self, mock_get_handler, mock_get_commit captured_push_shas = [] def capture_and_record( - repo_path, commit_sha, branch, session=None, issue_number=None, pipeline_phase=None, push_sha=None + repo_path, + commit_sha, + branch, + session=None, + issue_number=None, + pipeline_phase=None, + push_sha=None, ): captured_push_shas.append(push_sha) - return Checkpoint( + now = datetime.now(UTC) + return CheckpointV2( id=f"ckpt-{commit_sha[:12]}", + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, - session=SessionMetadata(session_id="test", started_at=datetime.now(UTC)), - created_at=datetime.now(UTC), + session_id="test", + session=SessionMetadata(session_id="test", started_at=now), + created_at=now, + session_started_at=now, push_sha=push_sha, branch=branch, ) mock_handler = MagicMock() mock_handler.capture_checkpoint.side_effect = capture_and_record - mock_handler.store_checkpoint.return_value = True + mock_handler.store_checkpoint_v2.return_value = True mock_get_handler.return_value = mock_handler import checkpoint_handler - checkpoints = checkpoint_handler.capture_and_store_checkpoints_for_push( + checkpoint_handler.capture_and_store_checkpoints_for_push( repo_path="/repo", old_sha="0" * 40, new_sha=commit2, @@ -270,7 +288,7 @@ def capture_and_record( @patch("checkpoint_handler.get_checkpoint_handler") def test_continues_on_individual_checkpoint_failure(self, mock_get_handler, mock_get_commits): """Test that failure to capture one checkpoint doesn't stop others.""" - from egg_contracts.checkpoints import Checkpoint, SessionMetadata + from egg_contracts.checkpoints import CheckpointV2, SessionMetadata, TriggerType # Mock 3 commits commit1 = "1111111111111111111111111111111111111111" @@ -281,23 +299,33 @@ def test_continues_on_individual_checkpoint_failure(self, mock_get_handler, mock call_count = [0] def capture_with_failure( - repo_path, commit_sha, branch, session=None, issue_number=None, pipeline_phase=None, push_sha=None + repo_path, + commit_sha, + branch, + session=None, + issue_number=None, + pipeline_phase=None, + push_sha=None, ): call_count[0] += 1 if call_count[0] == 2: raise Exception("Simulated failure") - return Checkpoint( + now = datetime.now(UTC) + return CheckpointV2( id=f"ckpt-{commit_sha[:12]}", + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, - session=SessionMetadata(session_id="test", started_at=datetime.now(UTC)), - created_at=datetime.now(UTC), + session_id="test", + session=SessionMetadata(session_id="test", started_at=now), + created_at=now, + session_started_at=now, push_sha=push_sha, branch=branch, ) mock_handler = MagicMock() mock_handler.capture_checkpoint.side_effect = capture_with_failure - mock_handler.store_checkpoint.return_value = True + mock_handler.store_checkpoint_v2.return_value = True mock_get_handler.return_value = mock_handler import checkpoint_handler @@ -313,23 +341,27 @@ def capture_with_failure( # Should have 2 checkpoints (first and third succeeded) assert len(checkpoints) == 2 assert call_count[0] == 3 - assert mock_handler.store_checkpoint.call_count == 2 + assert mock_handler.store_checkpoint_v2.call_count == 2 @patch("checkpoint_handler.get_commits_in_push") @patch("checkpoint_handler.get_checkpoint_handler") def test_async_store_uses_thread(self, mock_get_handler, mock_get_commits): """Test that async_store=True stores in a background thread.""" - from egg_contracts.checkpoints import Checkpoint, SessionMetadata + from egg_contracts.checkpoints import CheckpointV2, SessionMetadata, TriggerType commit1 = "1111111111111111111111111111111111111111" mock_get_commits.return_value = [commit1] mock_handler = MagicMock() - mock_handler.capture_checkpoint.return_value = Checkpoint( + now = datetime.now(UTC) + mock_handler.capture_checkpoint.return_value = CheckpointV2( id="ckpt-1111111111111", + trigger_type=TriggerType.COMMIT, commit_sha=commit1, - session=SessionMetadata(session_id="test", started_at=datetime.now(UTC)), - created_at=datetime.now(UTC), + session_id="test", + session=SessionMetadata(session_id="test", started_at=now), + created_at=now, + session_started_at=now, branch="main", ) mock_get_handler.return_value = mock_handler @@ -340,7 +372,7 @@ def test_async_store_uses_thread(self, mock_get_handler, mock_get_commits): mock_thread_instance = MagicMock() mock_thread.return_value = mock_thread_instance - checkpoints = checkpoint_handler.capture_and_store_checkpoints_for_push( + checkpoint_handler.capture_and_store_checkpoints_for_push( repo_path="/repo", old_sha="0" * 40, new_sha=commit1, @@ -353,13 +385,13 @@ def test_async_store_uses_thread(self, mock_get_handler, mock_get_commits): mock_thread_instance.start.assert_called_once() # store_checkpoint should NOT have been called directly (it's in the thread) - mock_handler.store_checkpoint.assert_not_called() + mock_handler.store_checkpoint_v2.assert_not_called() @patch("checkpoint_handler.get_commits_in_push") @patch("checkpoint_handler.get_checkpoint_handler") def test_checkpoint_returns_none_excluded(self, mock_get_handler, mock_get_commits): """Test that checkpoints returning None are excluded from result.""" - from egg_contracts.checkpoints import Checkpoint, SessionMetadata + from egg_contracts.checkpoints import CheckpointV2, SessionMetadata, TriggerType commit1 = "1111111111111111111111111111111111111111" commit2 = "2222222222222222222222222222222222222222" @@ -368,23 +400,33 @@ def test_checkpoint_returns_none_excluded(self, mock_get_handler, mock_get_commi call_count = [0] def sometimes_returns_none( - repo_path, commit_sha, branch, session=None, issue_number=None, pipeline_phase=None, push_sha=None + repo_path, + commit_sha, + branch, + session=None, + issue_number=None, + pipeline_phase=None, + push_sha=None, ): call_count[0] += 1 if call_count[0] == 1: return None - return Checkpoint( + now = datetime.now(UTC) + return CheckpointV2( id=f"ckpt-{commit_sha[:12]}", + trigger_type=TriggerType.COMMIT, commit_sha=commit_sha, - session=SessionMetadata(session_id="test", started_at=datetime.now(UTC)), - created_at=datetime.now(UTC), + session_id="test", + session=SessionMetadata(session_id="test", started_at=now), + created_at=now, + session_started_at=now, push_sha=push_sha, branch=branch, ) mock_handler = MagicMock() mock_handler.capture_checkpoint.side_effect = sometimes_returns_none - mock_handler.store_checkpoint.return_value = True + mock_handler.store_checkpoint_v2.return_value = True mock_get_handler.return_value = mock_handler import checkpoint_handler @@ -400,3 +442,253 @@ def sometimes_returns_none( # Only second commit's checkpoint should be returned assert len(checkpoints) == 1 assert checkpoints[0].commit_sha == commit2 + + +def _make_test_session( + container_id="test-container", + agent_role="coder", + phase="implement", + issue_number=530, + pr_number=42, +): + """Create a test Session for checkpoint tests.""" + now = datetime.now(UTC) + return Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id=container_id, + container_ip="172.18.0.5", + mode="private", + created_at=now - timedelta(hours=1), + last_seen=now, + expires_at=now + timedelta(hours=23), + agent_role=agent_role, + phase=phase, + issue_number=issue_number, + pr_number=pr_number, + ) + + +class TestCaptureSessionEndCheckpoint: + """Tests for capture_session_end_checkpoint function.""" + + def test_returns_none_when_disabled(self): + """Returns (None, None) when checkpoints are disabled.""" + import checkpoint_handler + + original = checkpoint_handler.CHECKPOINT_ENABLED + checkpoint_handler.CHECKPOINT_ENABLED = False + + try: + from egg_contracts.checkpoints import SessionStatus + + session = _make_test_session() + result = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.COMPLETED, + ) + assert result == (None, None) + finally: + checkpoint_handler.CHECKPOINT_ENABLED = original + + @patch("checkpoint_handler.get_checkpoint_handler") + def test_completed_session_creates_checkpoint(self, mock_get_handler): + """Session-end with COMPLETED status creates a checkpoint.""" + from egg_contracts.checkpoints import ( + CheckpointV2, + SessionMetadata, + SessionStatus, + TriggerType, + ) + + now = datetime.now(UTC) + mock_handler = MagicMock() + mock_handler.capture_session_end_checkpoint.return_value = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.COMPLETED, + session_id="test-container", + session=SessionMetadata(session_id="test-container", started_at=now), + created_at=now, + session_started_at=now, + ) + mock_handler.store_checkpoint_v2.return_value = True + mock_get_handler.return_value = mock_handler + + session = _make_test_session() + checkpoint, event = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.COMPLETED, + repo_path="/home/egg/repos/test-repo", + async_store=False, + ) + + assert checkpoint is not None + assert checkpoint.trigger_type == TriggerType.SESSION_END + assert checkpoint.session_status == SessionStatus.COMPLETED + assert event is None # sync store + + @patch("checkpoint_handler.get_checkpoint_handler") + def test_expired_session_creates_checkpoint(self, mock_get_handler): + """Session-end with EXPIRED status creates a checkpoint.""" + from egg_contracts.checkpoints import ( + CheckpointV2, + SessionMetadata, + SessionStatus, + TriggerType, + ) + + now = datetime.now(UTC) + mock_handler = MagicMock() + mock_handler.capture_session_end_checkpoint.return_value = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.EXPIRED, + session_id="test-container", + session=SessionMetadata(session_id="test-container", started_at=now), + created_at=now, + session_started_at=now, + ) + mock_get_handler.return_value = mock_handler + + session = _make_test_session() + checkpoint, event = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.EXPIRED, + repo_path="/home/egg/repos/test-repo", + async_store=False, + ) + + assert checkpoint is not None + assert checkpoint.session_status == SessionStatus.EXPIRED + + @patch("checkpoint_handler.get_checkpoint_handler") + def test_failed_session_creates_checkpoint(self, mock_get_handler): + """Session-end with FAILED status creates a checkpoint.""" + from egg_contracts.checkpoints import ( + CheckpointV2, + SessionMetadata, + SessionStatus, + TriggerType, + ) + + now = datetime.now(UTC) + mock_handler = MagicMock() + mock_handler.capture_session_end_checkpoint.return_value = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.FAILED, + session_id="test-container", + session=SessionMetadata(session_id="test-container", started_at=now), + created_at=now, + session_started_at=now, + ) + mock_get_handler.return_value = mock_handler + + session = _make_test_session() + checkpoint, event = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.FAILED, + repo_path="/home/egg/repos/test-repo", + async_store=False, + ) + + assert checkpoint is not None + assert checkpoint.session_status == SessionStatus.FAILED + + @patch("checkpoint_handler.get_checkpoint_handler") + def test_async_store_returns_completion_event(self, mock_get_handler): + """Async store returns a completion event that is eventually set.""" + from egg_contracts.checkpoints import ( + CheckpointV2, + SessionMetadata, + SessionStatus, + TriggerType, + ) + + now = datetime.now(UTC) + mock_handler = MagicMock() + mock_handler.capture_session_end_checkpoint.return_value = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.COMPLETED, + session_id="test-container", + session=SessionMetadata(session_id="test-container", started_at=now), + created_at=now, + session_started_at=now, + ) + mock_handler.store_checkpoint_v2.return_value = True + mock_get_handler.return_value = mock_handler + + session = _make_test_session() + checkpoint, event = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.COMPLETED, + repo_path="/home/egg/repos/test-repo", + async_store=True, + ) + + assert checkpoint is not None + assert event is not None + # Wait for async storage to complete + event.wait(timeout=5) + assert event.is_set() + mock_handler.store_checkpoint_v2.assert_called_once() + + @patch("checkpoint_handler.get_checkpoint_handler") + def test_capture_failure_returns_none(self, mock_get_handler): + """Returns (None, None) when handler fails to capture.""" + from egg_contracts.checkpoints import SessionStatus + + mock_handler = MagicMock() + mock_handler.capture_session_end_checkpoint.return_value = None + mock_get_handler.return_value = mock_handler + + session = _make_test_session() + checkpoint, event = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.COMPLETED, + ) + + assert checkpoint is None + assert event is None + + @patch("checkpoint_handler.get_checkpoint_handler") + def test_no_repo_path_returns_checkpoint_without_event(self, mock_get_handler): + """When no repo_path is available, returns checkpoint but no event.""" + from egg_contracts.checkpoints import ( + CheckpointV2, + SessionMetadata, + SessionStatus, + TriggerType, + ) + + now = datetime.now(UTC) + mock_handler = MagicMock() + mock_handler.capture_session_end_checkpoint.return_value = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.COMPLETED, + session_id="test-container", + session=SessionMetadata(session_id="test-container", started_at=now), + created_at=now, + session_started_at=now, + ) + mock_get_handler.return_value = mock_handler + + session = _make_test_session() + + # Patch Path to make repos dir not exist + with patch("checkpoint_handler.Path") as mock_path: + mock_repos_base = MagicMock() + mock_repos_base.exists.return_value = False + mock_path.return_value = mock_repos_base + + checkpoint, event = capture_session_end_checkpoint( + session=session, + session_status=SessionStatus.COMPLETED, + repo_path=None, + ) + + assert checkpoint is not None + assert event is None diff --git a/gateway/tests/test_policy.py b/gateway/tests/test_policy.py index 963a54ab4b..51d33be80e 100644 --- a/gateway/tests/test_policy.py +++ b/gateway/tests/test_policy.py @@ -1142,7 +1142,9 @@ def test_branch_push_blocked_reviewer_mode(self, policy_engine): def test_branch_push_blocked_reviewer_mode_even_bot_prefix(self, policy_engine): """Push to bot-prefixed branch is still blocked in reviewer mode.""" - result = policy_engine.check_branch_ownership("owner/repo", "egg-feature", auth_mode="reviewer") + result = policy_engine.check_branch_ownership( + "owner/repo", "egg-feature", auth_mode="reviewer" + ) assert not result.allowed assert "reviewer mode" in result.reason.lower() diff --git a/gateway/tests/test_session_manager.py b/gateway/tests/test_session_manager.py index 3e04e3c1a5..67e821c581 100644 --- a/gateway/tests/test_session_manager.py +++ b/gateway/tests/test_session_manager.py @@ -11,6 +11,7 @@ import pytest # Import from conftest-loaded module +import session_manager as session_manager_module from session_manager import ( Session, SessionManager, @@ -1105,3 +1106,321 @@ def test_from_persistence_handles_missing_phase(self): session = Session.from_persistence(data) assert session.phase is None + + +class TestSessionMetadataFields: + """Tests for issue_number and pr_number fields on Session.""" + + def test_session_with_issue_and_pr(self): + """Session can be created with issue_number and pr_number.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + issue_number=530, + pr_number=42, + ) + assert session.issue_number == 530 + assert session.pr_number == 42 + + def test_session_default_metadata_none(self): + """issue_number and pr_number default to None.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + ) + assert session.issue_number is None + assert session.pr_number is None + + def test_to_dict_includes_metadata(self): + """to_dict_for_persistence includes issue_number and pr_number.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + issue_number=530, + pr_number=42, + ) + d = session.to_dict_for_persistence() + assert d["issue_number"] == 530 + assert d["pr_number"] == 42 + + def test_to_dict_excludes_none_metadata(self): + """to_dict_for_persistence excludes None issue_number/pr_number.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + ) + d = session.to_dict_for_persistence() + assert "issue_number" not in d + assert "pr_number" not in d + + def test_roundtrip_with_metadata(self): + """issue_number and pr_number survive serialization roundtrip.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + issue_number=530, + pr_number=42, + ) + d = session.to_dict_for_persistence() + restored = Session.from_persistence(d) + + assert restored.issue_number == 530 + assert restored.pr_number == 42 + + def test_backward_compatibility_without_metadata(self): + """from_persistence handles sessions without issue_number/pr_number.""" + now = datetime.now(UTC) + data = { + "session_token_hash": "abc123", + "container_id": "test-container", + "container_ip": "172.18.0.5", + "mode": "private", + "created_at": now.isoformat(), + "last_seen": now.isoformat(), + "expires_at": (now + timedelta(hours=24)).isoformat(), + # No issue_number or pr_number - simulates legacy session + } + + session = Session.from_persistence(data) + assert session.issue_number is None + assert session.pr_number is None + + def test_register_session_with_metadata(self, tmp_path): + """register_session passes issue_number and pr_number.""" + manager = SessionManager(persistence_file=tmp_path / "sessions.json") + token, session = manager.register_session( + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + issue_number=530, + pr_number=42, + ) + assert session.issue_number == 530 + assert session.pr_number == 42 + + def test_metadata_persists_through_restart(self, tmp_path): + """issue_number and pr_number persist across gateway restarts.""" + persist_path = tmp_path / "sessions.json" + + manager1 = SessionManager(persistence_file=persist_path) + token, _ = manager1.register_session( + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + issue_number=530, + pr_number=42, + ) + + # Simulate restart + manager2 = SessionManager(persistence_file=persist_path) + session = manager2.get_session_by_container("test-container") + + assert session is not None + assert session.issue_number == 530 + assert session.pr_number == 42 + + +class TestSessionEndCheckpointCapture: + """Tests for session-end checkpoint capture during deletion/expiry.""" + + @pytest.fixture(autouse=True) + def clear_captured_containers(self): + """Clear the capture dedup set before each test.""" + session_manager_module._captured_containers.clear() + + @pytest.fixture + def manager(self, tmp_path): + """Create a session manager with a temporary persistence file.""" + return SessionManager(persistence_file=tmp_path / "sessions.json") + + def test_delete_session_by_token_captures_checkpoint(self, manager): + """delete_session(token) captures a session-end checkpoint.""" + token, session = manager.register_session( + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + ) + + from unittest.mock import patch + + with patch.object(session_manager_module, "_capture_and_cleanup_session") as mock_capture: + result = manager.delete_session(token) + assert result is True + mock_capture.assert_called_once() + args = mock_capture.call_args[0] + assert args[0].container_id == "test-container" + assert args[1] == "completed" + + def test_delete_session_by_container_captures_checkpoint(self, manager): + """delete_session_by_container captures a session-end checkpoint.""" + _token, session = manager.register_session( + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + ) + + from unittest.mock import patch + + with patch.object(session_manager_module, "_capture_and_cleanup_session") as mock_capture: + result = manager.delete_session_by_container("test-container") + assert result is True + mock_capture.assert_called_once() + args = mock_capture.call_args[0] + assert args[0].container_id == "test-container" + assert args[1] == "completed" + + def test_prune_captures_expired_checkpoints(self, manager): + """prune_expired_sessions captures checkpoints with EXPIRED status.""" + # Create and expire sessions + for i in range(3): + _token, session = manager.register_session( + container_id=f"expired-{i}", + container_ip="172.18.0.5", + mode="private", + ) + session.expires_at = datetime.now(UTC) - timedelta(seconds=1) + + from unittest.mock import patch + + with patch.object(session_manager_module, "_capture_and_cleanup_session") as mock_capture: + pruned = manager.prune_expired_sessions() + assert pruned == 3 + assert mock_capture.call_count == 3 + # All calls should use "expired" status + for c in mock_capture.call_args_list: + assert c[0][1] == "expired" + + def test_delete_by_token_not_found_skips_capture(self, manager): + """delete_session with invalid token doesn't capture checkpoint.""" + from unittest.mock import patch + + with patch.object(session_manager_module, "_capture_and_cleanup_session") as mock_capture: + result = manager.delete_session("nonexistent-token") + assert result is False + mock_capture.assert_not_called() + + def test_delete_by_container_not_found_skips_capture(self, manager): + """delete_session_by_container with invalid container doesn't capture.""" + from unittest.mock import patch + + with patch.object(session_manager_module, "_capture_and_cleanup_session") as mock_capture: + result = manager.delete_session_by_container("nonexistent") + assert result is False + mock_capture.assert_not_called() + + def test_capture_called_outside_lock(self, manager): + """Verify that checkpoint capture doesn't hold the session lock.""" + token, session = manager.register_session( + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + ) + + # Register another session to test lock contention + token2, _ = manager.register_session( + container_id="other-container", + container_ip="172.18.0.6", + mode="private", + ) + + from unittest.mock import patch + + def capture_mock(session_obj, status): + # While capture is running, other session operations should work + # If the lock were held, this would deadlock + result = manager.validate_session(token2) + assert result.valid + + with patch.object( + session_manager_module, "_capture_and_cleanup_session", side_effect=capture_mock + ): + result = manager.delete_session_by_container("test-container") + assert result is True + + def test_capture_and_cleanup_handles_import_error(self): + """_capture_and_cleanup_session handles ImportError gracefully.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + ) + + from unittest.mock import patch + + from session_manager import _capture_and_cleanup_session + + with patch.object(session_manager_module, "_cleanup_transcript_buffer") as mock_cleanup: + with patch.dict("sys.modules", {"checkpoint_handler": None}): + # Should not raise - just log and clean up + _capture_and_cleanup_session(session, "completed") + # Buffer cleanup should always happen + mock_cleanup.assert_called_once_with("test-container") + + def test_capture_and_cleanup_handles_capture_failure(self): + """_capture_and_cleanup_session cleans up buffer even on failure.""" + now = datetime.now(UTC) + session = Session( + session_token="test-token", + session_token_hash=_hash_token("test-token"), + container_id="test-container", + container_ip="172.18.0.5", + mode="private", + created_at=now, + last_seen=now, + expires_at=now + timedelta(hours=24), + ) + + from unittest.mock import patch + + from session_manager import _capture_and_cleanup_session + + with patch.object(session_manager_module, "_cleanup_transcript_buffer") as mock_cleanup: + with patch( + "checkpoint_handler.capture_session_end_checkpoint", + side_effect=RuntimeError("capture failed"), + ): + # Should not raise + _capture_and_cleanup_session(session, "completed") + # Buffer cleanup should still happen + mock_cleanup.assert_called_once_with("test-container") diff --git a/gateway/tests/test_transcript_buffer.py b/gateway/tests/test_transcript_buffer.py index 0f0e0c7837..9cf56d9ba7 100644 --- a/gateway/tests/test_transcript_buffer.py +++ b/gateway/tests/test_transcript_buffer.py @@ -1,21 +1,14 @@ """Tests for transcript_buffer module - API traffic capture.""" -import json -import tempfile import threading -import time -from pathlib import Path -from unittest.mock import patch, MagicMock - -import pytest +from unittest.mock import patch from transcript_buffer import ( + MAX_MESSAGE_CONTENT_LENGTH, TranscriptBuffer, - get_transcript_buffer, cleanup_transcript_buffer, get_buffer_path, - MAX_BUFFER_SIZE, - MAX_MESSAGE_CONTENT_LENGTH, + get_transcript_buffer, ) @@ -88,7 +81,10 @@ def test_write_multiple_api_turns(self, tmp_path): for i in range(3): buffer.write_api_turn( - request_body={"model": "claude", "messages": [{"role": "user", "content": f"Message {i}"}]}, + request_body={ + "model": "claude", + "messages": [{"role": "user", "content": f"Message {i}"}], + }, response_content=[{"type": "text", "text": f"Response {i}"}], response_usage={"input_tokens": 100 + i, "output_tokens": 50 + i}, streaming=False, @@ -126,8 +122,18 @@ def test_truncate_tools_to_names(self, tmp_path): "model": "claude", "messages": [], "tools": [ - {"name": "Bash", "type": "function", "description": "Run bash commands", "input_schema": {}}, - {"name": "Read", "type": "function", "description": "Read files", "input_schema": {}}, + { + "name": "Bash", + "type": "function", + "description": "Run bash commands", + "input_schema": {}, + }, + { + "name": "Read", + "type": "function", + "description": "Read files", + "input_schema": {}, + }, ], }, response_content=None, @@ -224,7 +230,10 @@ def test_rotation_when_size_exceeded(self, tmp_path): # Write entries until we exceed the limit for i in range(20): buffer.write_api_turn( - request_body={"model": "claude", "messages": [{"role": "user", "content": f"Message {i} with padding"}]}, + request_body={ + "model": "claude", + "messages": [{"role": "user", "content": f"Message {i} with padding"}], + }, response_content=[{"type": "text", "text": f"Response {i} with some padding text"}], response_usage={"input_tokens": 100, "output_tokens": 50}, streaming=False, @@ -248,7 +257,10 @@ def test_rotation_keeps_newest_entries(self, tmp_path): # Write entries with identifiable content for i in range(10): buffer.write_api_turn( - request_body={"model": "claude", "messages": [{"role": "user", "content": f"Message-{i}"}]}, + request_body={ + "model": "claude", + "messages": [{"role": "user", "content": f"Message-{i}"}], + }, response_content=[{"type": "text", "text": f"Response-{i}"}], response_usage={"input_tokens": i}, streaming=False, @@ -277,7 +289,10 @@ def test_concurrent_writes(self, tmp_path): def writer(thread_id): for i in range(writes_per_thread): success = buffer.write_api_turn( - request_body={"model": "claude", "messages": [{"role": "user", "content": f"Thread-{thread_id}-{i}"}]}, + request_body={ + "model": "claude", + "messages": [{"role": "user", "content": f"Thread-{thread_id}-{i}"}], + }, response_content=[{"type": "text", "text": f"Response-{thread_id}-{i}"}], response_usage={"thread": thread_id, "write": i}, streaming=False, diff --git a/gateway/tests/test_worktree_manager.py b/gateway/tests/test_worktree_manager.py index 45c00d29c7..a34e993c57 100644 --- a/gateway/tests/test_worktree_manager.py +++ b/gateway/tests/test_worktree_manager.py @@ -268,7 +268,9 @@ def test_with_active_containers(self, tmp_path): removed = startup_cleanup(active_containers={"active-container"}) assert removed == 1 - mock_instance.cleanup_orphaned_worktrees.assert_called_once_with({"active-container"}) + mock_instance.cleanup_orphaned_worktrees.assert_called_once_with( + {"active-container"}, None + ) def test_with_none_uses_docker(self): """Falls back to querying Docker when active_containers is None.""" @@ -388,7 +390,9 @@ def git_repo(self, tmp_path): repos_base.mkdir() repo_dir = repos_base / "test-repo" repo_dir.mkdir() - subprocess.run(["git", "init"], cwd=repo_dir, capture_output=True, check=True) + result = subprocess.run(["git", "init"], cwd=repo_dir, capture_output=True, text=True) + if result.returncode != 0: + pytest.skip(f"git init not available: {result.stderr.strip()}") subprocess.run( ["git", "commit", "--allow-empty", "-m", "init"], cwd=repo_dir, @@ -533,7 +537,6 @@ def test_no_false_match_without_git_suffix(self, tmp_path): "The fix ensures gitdir content must include /.git suffix to match." ) - def test_original_bug_comparison_without_git_suffix(self, tmp_path): """Demonstrate why comparing against worktree_path (not worktree_path/.git) was wrong. @@ -576,9 +579,7 @@ def test_original_bug_comparison_without_git_suffix(self, tmp_path): # NEW (fixed) comparison: gitdir_content == str(worktree_path / ".git") new_fixed_comparison = gitdir_content == str(pipeline_wt / ".git") - assert new_fixed_comparison, ( - "New comparison SHOULD match (both have /.git suffix)" - ) + assert new_fixed_comparison, "New comparison SHOULD match (both have /.git suffix)" # Verify the actual function returns the correct admin dir result = manager._find_worktree_git_dir(main_repo, pipeline_wt) diff --git a/gateway/worktree_manager.py b/gateway/worktree_manager.py index f88b607cff..c8265625b6 100644 --- a/gateway/worktree_manager.py +++ b/gateway/worktree_manager.py @@ -655,15 +655,23 @@ def list_worktrees(self) -> list[dict[str, Any]]: return worktrees - def cleanup_orphaned_worktrees(self, active_containers: set[str]) -> int: + def cleanup_orphaned_worktrees( + self, + active_containers: set[str], + session_manager: Any | None = None, + ) -> int: """ Remove worktrees for containers that no longer exist. Called on gateway startup and periodically to clean up orphaned worktrees from crashed containers. + For orphaned containers with active sessions, captures a session-end + checkpoint with FAILED status before cleaning up transcript buffers. + Args: active_containers: Set of currently active container IDs + session_manager: Optional SessionManager for session-end checkpoints Returns: Number of worktrees removed @@ -688,6 +696,21 @@ def cleanup_orphaned_worktrees(self, active_containers: set[str]) -> int: container_id=container_id, ) + # Capture session-end checkpoint for crashed container + if session_manager is not None: + try: + session = session_manager.get_session_by_container(container_id) + if session: + from session_manager import _capture_and_cleanup_session + + _capture_and_cleanup_session(session, "failed") + except Exception as e: + logger.warning( + "Failed to capture checkpoint for orphaned container", + container_id=container_id, + error=str(e), + ) + # Remove each worktree for worktree in list(container_dir.iterdir()): if worktree.is_dir(): @@ -762,7 +785,10 @@ def get_active_docker_containers() -> set[str]: return set() -def startup_cleanup(active_containers: set[str] | None = None) -> int: +def startup_cleanup( + active_containers: set[str] | None = None, + session_manager: Any | None = None, +) -> int: """ Clean up orphaned worktrees on gateway startup. @@ -774,6 +800,7 @@ def startup_cleanup(active_containers: set[str] | None = None) -> int: preserved. Pass an empty set when no containers are active. When None, falls back to querying Docker (which may not be available inside the gateway container). + session_manager: Optional SessionManager for session-end checkpoints Returns: Number of orphaned worktrees removed @@ -787,7 +814,7 @@ def startup_cleanup(active_containers: set[str] | None = None) -> int: active_containers=len(active_containers), ) - removed = manager.cleanup_orphaned_worktrees(active_containers) + removed = manager.cleanup_orphaned_worktrees(active_containers, session_manager) if removed > 0: logger.info(f"Cleaned up {removed} orphaned worktree(s)") diff --git a/integration_tests/local_pipeline/test_local_pipeline.py b/integration_tests/local_pipeline/test_local_pipeline.py index 74d38e6773..dfcc510283 100644 --- a/integration_tests/local_pipeline/test_local_pipeline.py +++ b/integration_tests/local_pipeline/test_local_pipeline.py @@ -986,8 +986,7 @@ def test_reviewer_mixed_verdict_handled(self, orchestrator_url: str) -> None: # At least one reviewed phase should have processed the mixed verdict reviewed_phases = ["refine", "plan", "implement"] any_reviewed = any( - phases.get(p, {}).get("status") == "complete" - for p in reviewed_phases + phases.get(p, {}).get("status") == "complete" for p in reviewed_phases ) assert any_reviewed, "At least one reviewed phase should have completed" @@ -1037,9 +1036,7 @@ def test_max_review_cycles_zero_triggers_immediate_circuit_breaker( if phase_name in phases: # review_cycles should not exceed 0 (circuit breaker fired immediately) cycles = phases[phase_name].get("review_cycles", 0) - assert cycles <= 1, ( - f"Phase {phase_name} had {cycles} review cycles with max=0" - ) + assert cycles <= 1, f"Phase {phase_name} had {cycles} review cycles with max=0" finally: delete_pipeline(orchestrator_url, pipeline_id) @@ -1084,9 +1081,7 @@ def test_review_cycle_counter_accuracy(self, orchestrator_url: str) -> None: finally: delete_pipeline(orchestrator_url, pipeline_id) - def test_reviewer_verdict_files_isolated_per_pipeline( - self, local_pipeline_stack - ) -> None: + def test_reviewer_verdict_files_isolated_per_pipeline(self, local_pipeline_stack) -> None: """Concurrent pipelines have distinct verdict file paths. Verifies that reviewer verdict files are properly namespaced by @@ -1152,9 +1147,7 @@ def test_reviewer_verdict_files_isolated_per_pipeline( delete_pipeline(orchestrator_url, pipeline_id_1) delete_pipeline(orchestrator_url, pipeline_id_2) - def test_pipeline_completes_despite_all_reviewers_reject( - self, orchestrator_url: str - ) -> None: + def test_pipeline_completes_despite_all_reviewers_reject(self, orchestrator_url: str) -> None: """Pipeline completes via circuit breaker when all reviewers reject. This test verifies graceful degradation: even if all reviewers diff --git a/integration_tests/local_pipeline/test_unified_pipeline_behavior.py b/integration_tests/local_pipeline/test_unified_pipeline_behavior.py index 43793e5882..8c1fc06862 100644 --- a/integration_tests/local_pipeline/test_unified_pipeline_behavior.py +++ b/integration_tests/local_pipeline/test_unified_pipeline_behavior.py @@ -14,8 +14,6 @@ """ import json -import random -import subprocess import time from pathlib import Path @@ -59,9 +57,7 @@ def test_contract_file_uses_pipeline_id_key(self, local_pipeline_stack) -> None: class TestLocalPipelinePrefixedPaths: """Local pipelines use pipeline_id-prefixed file paths.""" - def test_multiple_concurrent_pipelines_have_distinct_paths( - self, local_pipeline_stack - ) -> None: + def test_multiple_concurrent_pipelines_have_distinct_paths(self, local_pipeline_stack) -> None: """Two concurrent local pipelines should have distinct draft/verdict paths.""" orchestrator_url = local_pipeline_stack.orchestrator_url repos_dir = local_pipeline_stack.repos_dir @@ -93,9 +89,9 @@ def test_multiple_concurrent_pipelines_have_distinct_paths( for draft in draft_files: name = draft.name # Should be prefixed with one of the pipeline IDs - assert ( - name.startswith(pipeline_id_1) or name.startswith(pipeline_id_2) - ), f"Draft {name} should be prefixed with pipeline ID" + assert name.startswith(pipeline_id_1) or name.startswith(pipeline_id_2), ( + f"Draft {name} should be prefixed with pipeline ID" + ) finally: delete_pipeline(orchestrator_url, pipeline_id_1) @@ -130,9 +126,7 @@ def test_contract_synced_set_after_creation(self, local_pipeline_stack) -> None: class TestPhaseBasedPushRestrictions: """Push restrictions are per-phase, not per-mode.""" - def test_gateway_session_in_local_refine_allows_state_push( - self, local_pipeline_stack - ) -> None: + def test_gateway_session_in_local_refine_allows_state_push(self, local_pipeline_stack) -> None: """Local session in refine phase can push .egg-state/ files.""" gateway_url = local_pipeline_stack.gateway_url launcher_secret = local_pipeline_stack.launcher_secret @@ -234,9 +228,7 @@ def test_pipeline_state_file_exists(self, local_pipeline_stack) -> None: class TestPhasePromptInstructions: """Agent prompts include correct instructions for unified behavior.""" - def test_local_pipeline_includes_contract_cli_instructions( - self, local_pipeline_stack - ) -> None: + def test_local_pipeline_includes_contract_cli_instructions(self, local_pipeline_stack) -> None: """Local pipeline contract should exist with proper structure.""" orchestrator_url = local_pipeline_stack.orchestrator_url repos_dir = local_pipeline_stack.repos_dir @@ -268,13 +260,10 @@ class TestCheckpointBranchBypass: def test_checkpoint_branch_constant_defined(self) -> None: """Verify checkpoint branch constant is used consistently.""" # This is a structural test - the constant should be defined - # in gateway.py and should be "egg/checkpoints/v1" - expected_branch = "egg/checkpoints/v1" + # in gateway.py and should include "egg/checkpoints/v2" # Read gateway.py to verify the constant - gateway_path = ( - Path(__file__).parent.parent.parent / "gateway" / "gateway.py" - ) + gateway_path = Path(__file__).parent.parent.parent / "gateway" / "gateway.py" if gateway_path.exists(): content = gateway_path.read_text() - assert 'CHECKPOINT_BRANCH = "egg/checkpoints/v1"' in content + assert '"egg/checkpoints/v2"' in content diff --git a/orchestrator/container_monitor.py b/orchestrator/container_monitor.py index 05f0af1cbb..9a836c3520 100644 --- a/orchestrator/container_monitor.py +++ b/orchestrator/container_monitor.py @@ -8,9 +8,10 @@ import sys import threading import time -from datetime import datetime, timedelta +from collections.abc import Callable +from datetime import datetime from pathlib import Path -from typing import Any, Callable +from typing import Any # Add shared directory to path for logging _shared_path = Path(__file__).parent.parent / "shared" @@ -26,7 +27,7 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] return logging.getLogger(name) -from docker_client import DockerClient, ContainerNotFoundError, get_docker_client +from docker_client import ContainerNotFoundError, DockerClient, get_docker_client from models import ContainerInfo, ContainerStatus logger = get_logger("orchestrator.monitor") diff --git a/orchestrator/decision_queue.py b/orchestrator/decision_queue.py index c3e953c948..1c6951955b 100644 --- a/orchestrator/decision_queue.py +++ b/orchestrator/decision_queue.py @@ -8,10 +8,10 @@ import sys import threading import time +from collections.abc import Callable from datetime import datetime, timedelta from pathlib import Path -from typing import Any, Callable -from uuid import uuid4 +from typing import Any # Add shared directory to path for logging _shared_path = Path(__file__).parent.parent / "shared" @@ -185,10 +185,7 @@ def get_pending_decisions(self) -> list[HITLDecision]: """ with self._lock: pipeline = self._load_pipeline() - return [ - d for d in pipeline.decisions - if d.status == DecisionStatus.PENDING - ] + return [d for d in pipeline.decisions if d.status == DecisionStatus.PENDING] def get_decision(self, decision_id: str) -> HITLDecision: """Get a specific decision. @@ -306,9 +303,7 @@ def check_timeouts(self) -> list[HITLDecision]: if decision.status != DecisionStatus.PENDING: continue - timeout_at = decision.created_at + timedelta( - seconds=decision.timeout_seconds - ) + timeout_at = decision.created_at + timedelta(seconds=decision.timeout_seconds) if now > timeout_at: decision.status = DecisionStatus.TIMEOUT decision.resolved_at = now @@ -359,9 +354,7 @@ def wait_for_decision( if decision.status == DecisionStatus.RESOLVED: return decision elif decision.status in (DecisionStatus.TIMEOUT, DecisionStatus.CANCELLED): - raise DecisionTimeoutError( - f"Decision {decision_id} was {decision.status.value}" - ) + raise DecisionTimeoutError(f"Decision {decision_id} was {decision.status.value}") time.sleep(poll_interval) diff --git a/orchestrator/decision_timeout.py b/orchestrator/decision_timeout.py index 3c6b1728c5..f96e73a944 100644 --- a/orchestrator/decision_timeout.py +++ b/orchestrator/decision_timeout.py @@ -8,9 +8,9 @@ import sys import threading import time -from datetime import datetime, timedelta +from collections.abc import Callable +from datetime import datetime from pathlib import Path -from typing import Any, Callable # Add shared directory to path for logging _shared_path = Path(__file__).parent.parent / "shared" @@ -143,10 +143,7 @@ def _check_decision_status( timeout = decision.timeout_seconds should_timeout = elapsed >= timeout - should_warn = ( - not should_timeout - and elapsed >= timeout * self.warning_threshold_percent - ) + should_warn = not should_timeout and elapsed >= timeout * self.warning_threshold_percent return should_warn, should_timeout @@ -181,9 +178,10 @@ def _check_pipeline(self, pipeline_id: str, repo_path: Path) -> None: # Issue warning (only once per decision) self._warned_decisions.add(decision_key) - remaining = decision.timeout_seconds - ( - datetime.utcnow() - decision.created_at - ).total_seconds() + remaining = ( + decision.timeout_seconds + - (datetime.utcnow() - decision.created_at).total_seconds() + ) logger.warning( "Decision approaching timeout", @@ -254,6 +252,7 @@ def create_notification_handler( Returns: Escalation handler function """ + def handler(pipeline_id: str, decision: HITLDecision) -> None: message = ( f"Decision timeout in pipeline {pipeline_id}:\n" @@ -297,6 +296,7 @@ def create_slack_handler( Returns: Escalation handler function """ + def handler(pipeline_id: str, decision: HITLDecision) -> None: try: import requests diff --git a/orchestrator/dispatch.py b/orchestrator/dispatch.py index 927b7ffaec..6f690c4627 100644 --- a/orchestrator/dispatch.py +++ b/orchestrator/dispatch.py @@ -27,25 +27,22 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] from egg_contracts import load_contract, save_contract from egg_contracts.agent_roles import AgentRole as ContractAgentRole from egg_contracts.orchestrator import ( - AgentResult, DispatchDecision, - Orchestrator as ContractOrchestrator, collect_handoff_data, create_orchestrator, format_dispatch_for_workflow, - get_dispatch_for_contract, save_agent_output, ) - +from egg_contracts.orchestrator import ( + Orchestrator as ContractOrchestrator, +) from models import ( AgentExecution, AgentExecutionStatus, AgentRole, Pipeline, - PipelinePhase, - PipelineStatus, ) -from state_store import StateStore, get_state_store +from state_store import get_state_store logger = get_logger("orchestrator.dispatch") @@ -139,10 +136,7 @@ def get_agents_to_run(self) -> list[AgentRole]: List of AgentRole values """ decision = self.get_next_dispatch() - return [ - map_contract_role_to_agent_role(role) - for role in decision.agents_to_run - ] + return [map_contract_role_to_agent_role(role) for role in decision.agents_to_run] def is_complete(self) -> bool: """Check if all agents have completed. diff --git a/orchestrator/events.py b/orchestrator/events.py index 0dbc128ab1..3ed91efe12 100644 --- a/orchestrator/events.py +++ b/orchestrator/events.py @@ -7,12 +7,13 @@ import sys import threading +from collections.abc import Callable from dataclasses import dataclass, field from datetime import datetime -from enum import Enum +from enum import StrEnum from pathlib import Path -from typing import Any, Callable -from queue import Queue, Empty +from queue import Empty, Queue +from typing import Any # Add shared directory to path _shared_path = Path(__file__).parent.parent / "shared" @@ -31,7 +32,7 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] logger = get_logger("orchestrator.events") -class EventType(str, Enum): +class EventType(StrEnum): """Types of pipeline events.""" # Pipeline lifecycle diff --git a/orchestrator/handoffs.py b/orchestrator/handoffs.py index bcb48d0a25..88f37d2407 100644 --- a/orchestrator/handoffs.py +++ b/orchestrator/handoffs.py @@ -27,14 +27,16 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] from egg_contracts.agent_roles import ( AgentRole as ContractAgentRole, +) +from egg_contracts.agent_roles import ( get_role_definition, ) from egg_contracts.orchestrator import ( - collect_handoff_data as contract_collect_handoff, load_agent_output, +) +from egg_contracts.orchestrator import ( save_agent_output as contract_save_output, ) - from models import AgentExecution, AgentExecutionStatus, AgentRole logger = get_logger("orchestrator.handoffs") @@ -237,10 +239,7 @@ def get_handoff_env_var( """ handoffs = collect_handoff_data(repo_path, target_role) - data = { - role: handoff.data - for role, handoff in handoffs.items() - } + data = {role: handoff.data for role, handoff in handoffs.items()} return json.dumps(data) @@ -257,10 +256,7 @@ def get_agent_dependencies(role: AgentRole) -> list[AgentRole]: contract_role = ROLE_MAP[role] role_def = get_role_definition(contract_role) - return [ - REVERSE_ROLE_MAP[dep_role] - for dep_role in role_def.dependencies - ] + return [REVERSE_ROLE_MAP[dep_role] for dep_role in role_def.dependencies] def get_agent_dependents(role: AgentRole) -> list[AgentRole]: @@ -343,17 +339,11 @@ def add_result( def get_successful(self) -> list[AgentExecution]: """Get all successful executions.""" - return [ - ex for ex in self.results.values() - if ex.status == AgentExecutionStatus.COMPLETE - ] + return [ex for ex in self.results.values() if ex.status == AgentExecutionStatus.COMPLETE] def get_failed(self) -> list[AgentExecution]: """Get all failed executions.""" - return [ - ex for ex in self.results.values() - if ex.status == AgentExecutionStatus.FAILED - ] + return [ex for ex in self.results.values() if ex.status == AgentExecutionStatus.FAILED] def all_succeeded(self) -> bool: """Check if all agents succeeded.""" @@ -361,11 +351,7 @@ def all_succeeded(self) -> bool: def get_commits(self) -> list[str]: """Get all commit SHAs from successful agents.""" - return [ - ex.commit - for ex in self.get_successful() - if ex.commit - ] + return [ex.commit for ex in self.get_successful() if ex.commit] def get_summary(self) -> dict[str, Any]: """Get a summary of all results.""" diff --git a/orchestrator/metrics.py b/orchestrator/metrics.py index 45b9176b8b..080c2bf542 100644 --- a/orchestrator/metrics.py +++ b/orchestrator/metrics.py @@ -10,10 +10,8 @@ import sys import threading -import time from dataclasses import dataclass, field -from datetime import datetime, timedelta -from enum import Enum +from datetime import datetime from pathlib import Path from typing import Any @@ -112,7 +110,7 @@ def get_sum(self) -> float: def get_bucket_counts(self) -> dict[float, int]: """Get counts per bucket.""" with self._lock: - counts = {b: 0 for b in self.buckets} + counts = dict.fromkeys(self.buckets, 0) counts[float("inf")] = 0 for obs in self._observations: diff --git a/orchestrator/multi_agent.py b/orchestrator/multi_agent.py index ae9fe76a65..f843f19442 100644 --- a/orchestrator/multi_agent.py +++ b/orchestrator/multi_agent.py @@ -7,9 +7,10 @@ import sys import threading +from collections.abc import Callable from datetime import datetime from pathlib import Path -from typing import Any, Callable +from typing import Any # Add shared directory to path for logging _shared_path = Path(__file__).parent.parent / "shared" @@ -35,7 +36,6 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] ContainerInfo, ContainerStatus, Pipeline, - PipelineStatus, ) from sandbox_template import SandboxTemplate, create_sandbox_config from state_store import get_state_store diff --git a/orchestrator/resilience.py b/orchestrator/resilience.py index a59b9b2f77..d7cb50b69a 100644 --- a/orchestrator/resilience.py +++ b/orchestrator/resilience.py @@ -8,11 +8,12 @@ import sys import threading import time -from datetime import datetime, timedelta +from collections.abc import Callable +from datetime import datetime from enum import StrEnum from functools import wraps from pathlib import Path -from typing import Any, Callable, TypeVar +from typing import Any, TypeVar # Add shared directory to path for logging _shared_path = Path(__file__).parent.parent / "shared" @@ -203,7 +204,7 @@ def get_delay(self, attempt: int) -> float: Returns: Delay in seconds """ - delay = self.base_delay * (self.exponential_base ** attempt) + delay = self.base_delay * (self.exponential_base**attempt) delay = min(delay, self.max_delay) # Add jitter @@ -236,9 +237,7 @@ def decorator(func: Callable[..., T]) -> Callable[..., T]: def wrapper(*args: Any, **kwargs: Any) -> T: # Check circuit breaker if circuit_breaker and not circuit_breaker.allow_request(): - raise CircuitBreakerError( - f"Circuit breaker {circuit_breaker.name} is open" - ) + raise CircuitBreakerError(f"Circuit breaker {circuit_breaker.name} is open") last_error: Exception | None = None @@ -356,9 +355,7 @@ async def retry_async( # Check circuit breaker if circuit_breaker and not circuit_breaker.allow_request(): - raise CircuitBreakerError( - f"Circuit breaker {circuit_breaker.name} is open" - ) + raise CircuitBreakerError(f"Circuit breaker {circuit_breaker.name} is open") last_error: Exception | None = None diff --git a/orchestrator/routes/containers.py b/orchestrator/routes/containers.py index 49861a5298..9fb2ddafe2 100644 --- a/orchestrator/routes/containers.py +++ b/orchestrator/routes/containers.py @@ -38,7 +38,7 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] InvalidContainerIdError, get_docker_client, ) -from models import AgentRole, ContainerStatus +from models import AgentRole from sandbox_template import SandboxTemplate, create_sandbox_config logger = get_logger("orchestrator.containers") diff --git a/orchestrator/routes/pipelines.py b/orchestrator/routes/pipelines.py index bfce571465..cbf1a84371 100644 --- a/orchestrator/routes/pipelines.py +++ b/orchestrator/routes/pipelines.py @@ -84,6 +84,7 @@ def validate_checks(checks: list) -> list[dict[str, str]]: # type: ignore[misc] if isinstance(c, dict) and "name" in c and "command" in c ] + pipelines_bp = Blueprint("pipelines", __name__, url_prefix="/api/v1/pipelines") diff --git a/orchestrator/routes/sdlc_tokens.py b/orchestrator/routes/sdlc_tokens.py index c5369a49a2..430e677c8b 100644 --- a/orchestrator/routes/sdlc_tokens.py +++ b/orchestrator/routes/sdlc_tokens.py @@ -206,7 +206,9 @@ def approve_phase() -> tuple[Response, int]: return _make_error("Missing token") if phase not in VALID_PHASES: - return _make_error(f"Invalid phase: {phase}. Must be one of: {', '.join(sorted(VALID_PHASES))}") + return _make_error( + f"Invalid phase: {phase}. Must be one of: {', '.join(sorted(VALID_PHASES))}" + ) if pipeline_id not in _token_store: return _make_error(f"No tokens found for pipeline {pipeline_id}", 404) diff --git a/orchestrator/sdlc_wordlist.py b/orchestrator/sdlc_wordlist.py index f7f3576a61..fb6584f3d5 100644 --- a/orchestrator/sdlc_wordlist.py +++ b/orchestrator/sdlc_wordlist.py @@ -7,38 +7,219 @@ WORD_LIST = [ # Animals - "BEAR", "BIRD", "CAT", "CRAB", "CROW", "DEER", "DOG", "DOVE", "DUCK", - "EAGLE", "EEL", "ELK", "FISH", "FOX", "FROG", "GOAT", "GOOSE", "HAWK", - "HARE", "HORSE", "LAMB", "LION", "MOLE", "MOOSE", "MOTH", "MOUSE", - "NEWT", "OWL", "PANDA", "PIG", "PONY", "RAM", "ROBIN", "SEAL", "SHARK", - "SHEEP", "SLUG", "SNAIL", "SNAKE", "STORK", "SWAN", "TIGER", "TOAD", - "TROUT", "VIPER", "WASP", "WHALE", "WOLF", "WORM", "ZEBRA", + "BEAR", + "BIRD", + "CAT", + "CRAB", + "CROW", + "DEER", + "DOG", + "DOVE", + "DUCK", + "EAGLE", + "EEL", + "ELK", + "FISH", + "FOX", + "FROG", + "GOAT", + "GOOSE", + "HAWK", + "HARE", + "HORSE", + "LAMB", + "LION", + "MOLE", + "MOOSE", + "MOTH", + "MOUSE", + "NEWT", + "OWL", + "PANDA", + "PIG", + "PONY", + "RAM", + "ROBIN", + "SEAL", + "SHARK", + "SHEEP", + "SLUG", + "SNAIL", + "SNAKE", + "STORK", + "SWAN", + "TIGER", + "TOAD", + "TROUT", + "VIPER", + "WASP", + "WHALE", + "WOLF", + "WORM", + "ZEBRA", # Colors - "AMBER", "BLACK", "BLUE", "BROWN", "CORAL", "CREAM", "GOLD", "GREEN", - "GREY", "IVORY", "LILAC", "NAVY", "OLIVE", "PEACH", "PINK", "PLUM", - "RED", "RUBY", "RUST", "TAN", "TEAL", "WHITE", + "AMBER", + "BLACK", + "BLUE", + "BROWN", + "CORAL", + "CREAM", + "GOLD", + "GREEN", + "GREY", + "IVORY", + "LILAC", + "NAVY", + "OLIVE", + "PEACH", + "PINK", + "PLUM", + "RED", + "RUBY", + "RUST", + "TAN", + "TEAL", + "WHITE", # Fruits & food - "APPLE", "BASIL", "BEAN", "BERRY", "BREAD", "CANDY", "CEDAR", "CHERRY", - "CHILI", "COCOA", "DATE", "FIG", "GRAPE", "GUAVA", "HONEY", "KIWI", - "LEMON", "LIME", "MANGO", "MAPLE", "MELON", "MINT", - "PEAR", "RICE", "SAGE", "WHEAT", + "APPLE", + "BASIL", + "BEAN", + "BERRY", + "BREAD", + "CANDY", + "CEDAR", + "CHERRY", + "CHILI", + "COCOA", + "DATE", + "FIG", + "GRAPE", + "GUAVA", + "HONEY", + "KIWI", + "LEMON", + "LIME", + "MANGO", + "MAPLE", + "MELON", + "MINT", + "PEAR", + "RICE", + "SAGE", + "WHEAT", # Weather & sky - "BOLT", "BREEZE", "CLOUD", "DEW", "FLAME", "FLOOD", "FOG", "FROST", - "GALE", "HAIL", "ICE", "MIST", "MOON", "RAIN", "SKY", "SLEET", - "SNOW", "STAR", "STORM", "SUN", "TIDE", "WIND", + "BOLT", + "BREEZE", + "CLOUD", + "DEW", + "FLAME", + "FLOOD", + "FOG", + "FROST", + "GALE", + "HAIL", + "ICE", + "MIST", + "MOON", + "RAIN", + "SKY", + "SLEET", + "SNOW", + "STAR", + "STORM", + "SUN", + "TIDE", + "WIND", # Nature & landscape - "ASH", "BAY", "BLUFF", "BOG", "BROOK", "CAPE", "CAVE", "CLIFF", - "COAST", "COVE", "CREEK", "DALE", "DELTA", "DUNE", "FALLS", "FIELD", - "FJORD", "GLEN", "GROVE", "GULF", "HILL", "ISLE", "LAKE", "MARSH", - "MESA", "MOSS", "OAK", "PALM", "PEAK", "PINE", "POND", "REEF", - "RIDGE", "RIVER", "ROCK", "SAND", "SHORE", "SLOPE", "STONE", "TRAIL", - "VALE", "VINE", "WOOD", + "ASH", + "BAY", + "BLUFF", + "BOG", + "BROOK", + "CAPE", + "CAVE", + "CLIFF", + "COAST", + "COVE", + "CREEK", + "DALE", + "DELTA", + "DUNE", + "FALLS", + "FIELD", + "FJORD", + "GLEN", + "GROVE", + "GULF", + "HILL", + "ISLE", + "LAKE", + "MARSH", + "MESA", + "MOSS", + "OAK", + "PALM", + "PEAK", + "PINE", + "POND", + "REEF", + "RIDGE", + "RIVER", + "ROCK", + "SAND", + "SHORE", + "SLOPE", + "STONE", + "TRAIL", + "VALE", + "VINE", + "WOOD", # Tools & objects - "ANVIL", "AXE", "BELL", "BOW", "BRICK", "BRUSH", "CHAIN", - "CHEST", "CLOCK", "COIN", "CROWN", "DRUM", "FLUTE", "FORGE", "GATE", - "GEAR", "GLASS", "GLOBE", "GONG", "HARP", "HELM", "HORN", "KEY", - "KNOT", "LANCE", "LATCH", "LENS", "LOCK", "MASK", "NAIL", "OPAL", - "PEARL", "PIKE", "PIPE", "PRISM", "RING", "ROPE", "SCALE", "SHIELD", - "SPEAR", "SPIKE", "SWORD", "TORCH", "TOWER", "VAULT", "WAGON", + "ANVIL", + "AXE", + "BELL", + "BOW", + "BRICK", + "BRUSH", + "CHAIN", + "CHEST", + "CLOCK", + "COIN", + "CROWN", + "DRUM", + "FLUTE", + "FORGE", + "GATE", + "GEAR", + "GLASS", + "GLOBE", + "GONG", + "HARP", + "HELM", + "HORN", + "KEY", + "KNOT", + "LANCE", + "LATCH", + "LENS", + "LOCK", + "MASK", + "NAIL", + "OPAL", + "PEARL", + "PIKE", + "PIPE", + "PRISM", + "RING", + "ROPE", + "SCALE", + "SHIELD", + "SPEAR", + "SPIKE", + "SWORD", + "TORCH", + "TOWER", + "VAULT", + "WAGON", "WHEEL", ] diff --git a/orchestrator/sse.py b/orchestrator/sse.py index d5e5d0c5cf..bf91701489 100644 --- a/orchestrator/sse.py +++ b/orchestrator/sse.py @@ -303,6 +303,7 @@ def create_sse_stream( if repo_path is None: try: from routes import get_repo_path + repo_path = get_repo_path() except Exception: pass @@ -371,9 +372,7 @@ def create_sse_stream( except Empty: # No events — send heartbeat to keep connection alive - yield format_sse_comment( - f"heartbeat {datetime.utcnow().isoformat()}Z" - ) + yield format_sse_comment(f"heartbeat {datetime.utcnow().isoformat()}Z") finally: manager.remove_client(pipeline_id, q) diff --git a/orchestrator/tests/test_cli.py b/orchestrator/tests/test_cli.py index 24bc45b2e8..5aa787fd22 100644 --- a/orchestrator/tests/test_cli.py +++ b/orchestrator/tests/test_cli.py @@ -8,7 +8,6 @@ from unittest.mock import MagicMock, patch import pytest - from cli import create_parser, main diff --git a/orchestrator/tests/test_dag_visualizer.py b/orchestrator/tests/test_dag_visualizer.py index 0c8a4d8d1f..0f56327110 100644 --- a/orchestrator/tests/test_dag_visualizer.py +++ b/orchestrator/tests/test_dag_visualizer.py @@ -202,9 +202,7 @@ def test_phase_with_agents(self): ], ) } - pipeline = create_test_pipeline( - phases=phases, current_phase=PipelinePhase.IMPLEMENT - ) + pipeline = create_test_pipeline(phases=phases, current_phase=PipelinePhase.IMPLEMENT) result = render_pipeline_dag(pipeline) assert "coder" in result @@ -237,9 +235,7 @@ def test_phase_with_mixed_agent_statuses(self): ], ) } - pipeline = create_test_pipeline( - phases=phases, current_phase=PipelinePhase.IMPLEMENT - ) + pipeline = create_test_pipeline(phases=phases, current_phase=PipelinePhase.IMPLEMENT) result = render_pipeline_dag(pipeline) assert "✓ coder" in result @@ -274,9 +270,7 @@ def test_phase_agents_ascii_mode(self): ], ) } - pipeline = create_test_pipeline( - phases=phases, current_phase=PipelinePhase.IMPLEMENT - ) + pipeline = create_test_pipeline(phases=phases, current_phase=PipelinePhase.IMPLEMENT) result = render_pipeline_dag(pipeline, use_ascii=True) assert "+ coder" in result diff --git a/orchestrator/tests/test_docker_client.py b/orchestrator/tests/test_docker_client.py index 8111882142..ac7674d653 100644 --- a/orchestrator/tests/test_docker_client.py +++ b/orchestrator/tests/test_docker_client.py @@ -9,12 +9,9 @@ from unittest.mock import MagicMock, patch import pytest - from docker_client import ( ContainerNotFoundError, - ContainerOperationError, DockerClient, - DockerClientError, ImageNotFoundError, InvalidContainerIdError, _validate_container_id, @@ -76,7 +73,6 @@ def test_create_container(self, docker_client, mock_docker): def test_create_container_image_not_found(self, docker_client, mock_docker): """Test create fails when image not found.""" - from docker.errors import ImageNotFound mock_docker.images.get.return_value = None diff --git a/orchestrator/tests/test_events.py b/orchestrator/tests/test_events.py index 89c12323b9..6141f3751c 100644 --- a/orchestrator/tests/test_events.py +++ b/orchestrator/tests/test_events.py @@ -4,10 +4,8 @@ import time from datetime import datetime -from unittest.mock import MagicMock import pytest - from events import ( Event, EventBus, @@ -235,6 +233,7 @@ def test_emit_event(self): """Test emit_event function.""" # Reset singleton import events + events._event_bus = None received_events = [] diff --git a/orchestrator/tests/test_metrics.py b/orchestrator/tests/test_metrics.py index c1b923a6b5..d50400051e 100644 --- a/orchestrator/tests/test_metrics.py +++ b/orchestrator/tests/test_metrics.py @@ -255,6 +255,7 @@ class TestSingletonMetrics: def test_get_metrics_registry_singleton(self): """Test registry singleton.""" import metrics + metrics._registry = None registry1 = get_metrics_registry() @@ -264,6 +265,7 @@ def test_get_metrics_registry_singleton(self): def test_get_metrics_singleton(self): """Test metrics singleton.""" import metrics + metrics._metrics = None metrics1 = get_metrics() diff --git a/orchestrator/tests/test_sdlc_tokens.py b/orchestrator/tests/test_sdlc_tokens.py index 973dab85c8..38efd5204d 100644 --- a/orchestrator/tests/test_sdlc_tokens.py +++ b/orchestrator/tests/test_sdlc_tokens.py @@ -361,8 +361,10 @@ def test_token_gated_pipeline_blocks_direct_resolution(self): mock_store.load_pipeline.return_value = mock_pipeline with test_app.test_client() as c: - with patch("routes.decisions.get_state_store", return_value=mock_store), \ - patch("routes.decisions.get_repo_path", return_value="/tmp/test"): + with ( + patch("routes.decisions.get_state_store", return_value=mock_store), + patch("routes.decisions.get_repo_path", return_value="/tmp/test"), + ): resp = c.post( "/api/v1/pipelines/issue-500/decisions/decision-1/resolve", json={"resolution": "approve"}, @@ -381,8 +383,10 @@ def test_state_store_exception_returns_503(self): mock_store.load_pipeline.side_effect = RuntimeError("connection lost") with test_app.test_client() as c: - with patch("routes.decisions.get_state_store", return_value=mock_store), \ - patch("routes.decisions.get_repo_path", return_value="/tmp/test"): + with ( + patch("routes.decisions.get_state_store", return_value=mock_store), + patch("routes.decisions.get_repo_path", return_value="/tmp/test"), + ): resp = c.post( "/api/v1/pipelines/issue-500/decisions/decision-1/resolve", json={"resolution": "approve"}, @@ -412,9 +416,11 @@ def test_pipeline_not_found_allows_resolution(self): mock_queue.resolve_decision.return_value = mock_decision with test_app.test_client() as c: - with patch("routes.decisions.get_state_store", return_value=mock_store), \ - patch("routes.decisions.get_repo_path", return_value="/tmp/test"), \ - patch("routes.decisions.get_decision_queue", return_value=mock_queue): + with ( + patch("routes.decisions.get_state_store", return_value=mock_store), + patch("routes.decisions.get_repo_path", return_value="/tmp/test"), + patch("routes.decisions.get_decision_queue", return_value=mock_queue), + ): resp = c.post( "/api/v1/pipelines/issue-500/decisions/decision-1/resolve", json={"resolution": "approve"}, @@ -431,7 +437,9 @@ def test_matches_exact_phase_transition_question(self): from routes.sdlc_tokens import _is_phase_transition_decision decision = MagicMock() - decision.question = "The refine phase has completed. Please review the analysis and approve to continue." + decision.question = ( + "The refine phase has completed. Please review the analysis and approve to continue." + ) assert _is_phase_transition_decision(decision, "refine") is True def test_does_not_match_different_phase(self): @@ -439,7 +447,9 @@ def test_does_not_match_different_phase(self): from routes.sdlc_tokens import _is_phase_transition_decision decision = MagicMock() - decision.question = "The refine phase has completed. Please review the analysis and approve to continue." + decision.question = ( + "The refine phase has completed. Please review the analysis and approve to continue." + ) assert _is_phase_transition_decision(decision, "plan") is False def test_does_not_match_unrelated_question_containing_phase_name(self): @@ -464,7 +474,9 @@ def test_matches_plan_phase_transition(self): from routes.sdlc_tokens import _is_phase_transition_decision decision = MagicMock() - decision.question = "The plan phase has completed. Please review the plan and approve to continue." + decision.question = ( + "The plan phase has completed. Please review the plan and approve to continue." + ) assert _is_phase_transition_decision(decision, "plan") is True @@ -478,13 +490,17 @@ def test_resolves_matching_decision(self): mock_decision = MagicMock() mock_decision.id = "decision-1" - mock_decision.question = "The refine phase has completed. Please review the analysis and approve to continue." + mock_decision.question = ( + "The refine phase has completed. Please review the analysis and approve to continue." + ) mock_queue = MagicMock() mock_queue.get_pending_decisions.return_value = [mock_decision] - with patch("routes.get_repo_path", return_value="/tmp/test"), \ - patch("decision_queue.get_decision_queue", return_value=mock_queue): + with ( + patch("routes.get_repo_path", return_value="/tmp/test"), + patch("decision_queue.get_decision_queue", return_value=mock_queue), + ): _resolve_phase_decisions("issue-600", "refine") mock_queue.resolve_decision.assert_called_once_with( @@ -502,8 +518,10 @@ def test_does_not_resolve_unrelated_decision(self): mock_queue = MagicMock() mock_queue.get_pending_decisions.return_value = [mock_decision] - with patch("routes.get_repo_path", return_value="/tmp/test"), \ - patch("decision_queue.get_decision_queue", return_value=mock_queue): + with ( + patch("routes.get_repo_path", return_value="/tmp/test"), + patch("decision_queue.get_decision_queue", return_value=mock_queue), + ): _resolve_phase_decisions("issue-600", "plan") mock_queue.resolve_decision.assert_not_called() @@ -537,8 +555,10 @@ def test_reset_clears_in_memory_tokens(self, client, auth_headers): mock_store = MagicMock() mock_store.load_pipeline.side_effect = PipelineNotFoundError("not found") - with patch("state_store.get_state_store", return_value=mock_store), \ - patch("routes.get_repo_path", return_value="/tmp/test"): + with ( + patch("state_store.get_state_store", return_value=mock_store), + patch("routes.get_repo_path", return_value="/tmp/test"), + ): resp = client.post( "/api/v1/sdlc-tokens/reset", json={"pipeline_id": "issue-700"}, @@ -563,8 +583,10 @@ def test_reset_nonexistent_pipeline(self, client, auth_headers): mock_store = MagicMock() mock_store.load_pipeline.side_effect = PipelineNotFoundError("not found") - with patch("state_store.get_state_store", return_value=mock_store), \ - patch("routes.get_repo_path", return_value="/tmp/test"): + with ( + patch("state_store.get_state_store", return_value=mock_store), + patch("routes.get_repo_path", return_value="/tmp/test"), + ): resp = client.post( "/api/v1/sdlc-tokens/reset", json={"pipeline_id": "issue-999"}, @@ -591,8 +613,10 @@ def test_reset_returns_503_on_persistent_store_failure(self, client, auth_header mock_store.load_pipeline.return_value = mock_pipeline mock_store.save_pipeline.side_effect = RuntimeError("disk full") - with patch("state_store.get_state_store", return_value=mock_store), \ - patch("routes.get_repo_path", return_value="/tmp/test"): + with ( + patch("state_store.get_state_store", return_value=mock_store), + patch("routes.get_repo_path", return_value="/tmp/test"), + ): resp = client.post( "/api/v1/sdlc-tokens/reset", json={"pipeline_id": "issue-750"}, @@ -620,8 +644,10 @@ def test_reset_clears_both_stores_on_success(self, client, auth_headers): mock_store = MagicMock() mock_store.load_pipeline.return_value = mock_pipeline - with patch("state_store.get_state_store", return_value=mock_store), \ - patch("routes.get_repo_path", return_value="/tmp/test"): + with ( + patch("state_store.get_state_store", return_value=mock_store), + patch("routes.get_repo_path", return_value="/tmp/test"), + ): resp = client.post( "/api/v1/sdlc-tokens/reset", json={"pipeline_id": "issue-760"}, diff --git a/orchestrator/tests/test_sse.py b/orchestrator/tests/test_sse.py index 0f4fdc3cb8..bf8f8a415b 100644 --- a/orchestrator/tests/test_sse.py +++ b/orchestrator/tests/test_sse.py @@ -243,8 +243,10 @@ def test_event_includes_visualization_when_available(self): q = manager.add_client("test-123") - with patch("sse.get_state_store") as mock_store, \ - patch("sse.render_pipeline_dag") as mock_dag: + with ( + patch("sse.get_state_store") as mock_store, + patch("sse.render_pipeline_dag") as mock_dag, + ): store_instance = MagicMock() store_instance.load_pipeline.return_value = pipeline mock_store.return_value = store_instance @@ -276,9 +278,11 @@ def test_event_visualization_falls_back_to_env_var(self): q = manager.add_client("test-123") - with patch("sse.get_state_store") as mock_store, \ - patch("sse.render_pipeline_dag") as mock_dag, \ - patch.dict("os.environ", {"EGG_REPO_PATH": "/tmp/test-env-repo"}): + with ( + patch("sse.get_state_store") as mock_store, + patch("sse.render_pipeline_dag") as mock_dag, + patch.dict("os.environ", {"EGG_REPO_PATH": "/tmp/test-env-repo"}), + ): store_instance = MagicMock() store_instance.load_pipeline.return_value = pipeline mock_store.return_value = store_instance @@ -307,11 +311,12 @@ def test_already_terminal_pipeline(self): pipeline = create_test_pipeline(status=PipelineStatus.COMPLETE) tmp_dir = Path("/tmp/test-sse-stream") - with patch("sse.get_sse_manager") as mock_mgr, \ - patch("sse.get_state_store") as mock_store, \ - patch("sse.render_pipeline_dag") as mock_dag, \ - patch("sse.generate_status_report") as mock_report: - + with ( + patch("sse.get_sse_manager") as mock_mgr, + patch("sse.get_state_store") as mock_store, + patch("sse.render_pipeline_dag") as mock_dag, + patch("sse.generate_status_report") as mock_report, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -339,9 +344,7 @@ def test_pipeline_not_found(self): tmp_dir = Path("/tmp/test-sse-notfound") - with patch("sse.get_sse_manager") as mock_mgr, \ - patch("sse.get_state_store") as mock_store: - + with patch("sse.get_sse_manager") as mock_mgr, patch("sse.get_state_store") as mock_store: mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -363,11 +366,12 @@ def test_cleanup_on_generator_close(self): pipeline = create_test_pipeline() tmp_dir = Path("/tmp/test-sse-cleanup") - with patch("sse.get_sse_manager") as mock_mgr, \ - patch("sse.get_state_store") as mock_store, \ - patch("sse.render_pipeline_dag") as mock_dag, \ - patch("sse.generate_status_report") as mock_report: - + with ( + patch("sse.get_sse_manager") as mock_mgr, + patch("sse.get_state_store") as mock_store, + patch("sse.render_pipeline_dag") as mock_dag, + patch("sse.generate_status_report") as mock_report, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -391,11 +395,12 @@ def test_snapshot_includes_visualization(self): pipeline = create_test_pipeline() tmp_dir = Path("/tmp/test-sse-snapshot") - with patch("sse.get_sse_manager") as mock_mgr, \ - patch("sse.get_state_store") as mock_store, \ - patch("sse.render_pipeline_dag") as mock_dag, \ - patch("sse.generate_status_report") as mock_report: - + with ( + patch("sse.get_sse_manager") as mock_mgr, + patch("sse.get_state_store") as mock_store, + patch("sse.render_pipeline_dag") as mock_dag, + patch("sse.generate_status_report") as mock_report, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -424,11 +429,12 @@ def test_terminal_event_ends_stream(self): pipeline = create_test_pipeline() tmp_dir = Path("/tmp/test-sse-terminal") - with patch("sse.get_sse_manager") as mock_mgr, \ - patch("sse.get_state_store") as mock_store, \ - patch("sse.render_pipeline_dag") as mock_dag, \ - patch("sse.generate_status_report") as mock_report: - + with ( + patch("sse.get_sse_manager") as mock_mgr, + patch("sse.get_state_store") as mock_store, + patch("sse.render_pipeline_dag") as mock_dag, + patch("sse.generate_status_report") as mock_report, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -461,7 +467,6 @@ def test_terminal_event_ends_stream(self): def test_no_initial_snapshot(self): """Test stream without initial snapshot.""" with patch("sse.get_sse_manager") as mock_mgr: - mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q diff --git a/orchestrator/tests/test_unified_sse.py b/orchestrator/tests/test_unified_sse.py index 9bbd355889..50dae7b454 100644 --- a/orchestrator/tests/test_unified_sse.py +++ b/orchestrator/tests/test_unified_sse.py @@ -124,16 +124,20 @@ def test_events_from_different_pipelines(self): q = manager.add_client() # Events from two different pipelines - bus.publish(Event( - event_type=EventType.PHASE_STARTED, - pipeline_id="pipeline-1", - data={"phase": "plan"}, - )) - bus.publish(Event( - event_type=EventType.PHASE_COMPLETED, - pipeline_id="pipeline-2", - data={"phase": "implement"}, - )) + bus.publish( + Event( + event_type=EventType.PHASE_STARTED, + pipeline_id="pipeline-1", + data={"phase": "plan"}, + ) + ) + bus.publish( + Event( + event_type=EventType.PHASE_COMPLETED, + pipeline_id="pipeline-2", + data={"phase": "implement"}, + ) + ) assert q.qsize() == 2 p1 = q.get_nowait() @@ -270,11 +274,12 @@ def test_initial_snapshot_with_active_pipelines(self): create_test_pipeline("p-2", PipelineStatus.RUNNING), ] - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._collect_all_pipelines_safe") as mock_collect, \ - patch("unified_sse.render_compact_status") as mock_compact, \ - patch("unified_sse.render_progress_bar") as mock_progress: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._collect_all_pipelines_safe") as mock_collect, + patch("unified_sse.render_compact_status") as mock_compact, + patch("unified_sse.render_progress_bar") as mock_progress, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -292,9 +297,7 @@ def test_initial_snapshot_with_active_pipelines(self): assert "retry: 5000" in snapshot # Parse the snapshot data - data_line = [ - l for l in snapshot.split("\n") if l.startswith("data: ") - ][0] + data_line = [l for l in snapshot.split("\n") if l.startswith("data: ")][0] data = json.loads(data_line[6:]) assert len(data["pipelines"]) == 2 assert data["pipelines"][0]["pipeline_id"] == "p-1" @@ -308,11 +311,12 @@ def test_snapshot_filters_terminal_pipelines(self): create_test_pipeline("p-3", PipelineStatus.FAILED), ] - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._collect_all_pipelines_safe") as mock_collect, \ - patch("unified_sse.render_compact_status") as mock_compact, \ - patch("unified_sse.render_progress_bar") as mock_progress: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._collect_all_pipelines_safe") as mock_collect, + patch("unified_sse.render_compact_status") as mock_compact, + patch("unified_sse.render_progress_bar") as mock_progress, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -322,15 +326,11 @@ def test_snapshot_filters_terminal_pipelines(self): mock_compact.return_value = "COMPACT" mock_progress.return_value = "[##] 50%" - gen = create_unified_sse_stream( - repo_path=Path("/tmp/test"), active_only=True - ) + gen = create_unified_sse_stream(repo_path=Path("/tmp/test"), active_only=True) snapshot = next(gen) gen.close() - data_line = [ - l for l in snapshot.split("\n") if l.startswith("data: ") - ][0] + data_line = [l for l in snapshot.split("\n") if l.startswith("data: ")][0] data = json.loads(data_line[6:]) assert len(data["pipelines"]) == 1 assert data["pipelines"][0]["pipeline_id"] == "p-1" @@ -342,11 +342,12 @@ def test_snapshot_includes_all_when_active_only_false(self): create_test_pipeline("p-2", PipelineStatus.COMPLETE), ] - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._collect_all_pipelines_safe") as mock_collect, \ - patch("unified_sse.render_compact_status") as mock_compact, \ - patch("unified_sse.render_progress_bar") as mock_progress: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._collect_all_pipelines_safe") as mock_collect, + patch("unified_sse.render_compact_status") as mock_compact, + patch("unified_sse.render_progress_bar") as mock_progress, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -356,23 +357,20 @@ def test_snapshot_includes_all_when_active_only_false(self): mock_compact.return_value = "COMPACT" mock_progress.return_value = "[##] 50%" - gen = create_unified_sse_stream( - repo_path=Path("/tmp/test"), active_only=False - ) + gen = create_unified_sse_stream(repo_path=Path("/tmp/test"), active_only=False) snapshot = next(gen) gen.close() - data_line = [ - l for l in snapshot.split("\n") if l.startswith("data: ") - ][0] + data_line = [l for l in snapshot.split("\n") if l.startswith("data: ")][0] data = json.loads(data_line[6:]) assert len(data["pipelines"]) == 2 def test_terminal_event_does_not_end_stream(self): """Test that terminal event for one pipeline doesn't end the unified stream.""" - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -409,12 +407,14 @@ def test_terminal_event_does_not_end_stream(self): def test_heartbeat_on_idle(self): """Test that heartbeat is sent when no events arrive.""" - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path, \ - patch("unified_sse.HEARTBEAT_INTERVAL", 0.1): - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + patch("unified_sse.HEARTBEAT_INTERVAL", 0.1), + ): # Need to also patch the imported constant import unified_sse + original_interval = unified_sse.HEARTBEAT_INTERVAL mock_q = Queue() @@ -439,11 +439,12 @@ def test_heartbeat_on_idle(self): def test_max_connection_time(self): """Test that stream ends after max connection time.""" - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path, \ - patch("unified_sse.MAX_CONNECTION_TIME", 0), \ - patch("unified_sse.HEARTBEAT_INTERVAL", 0.01): - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + patch("unified_sse.MAX_CONNECTION_TIME", 0), + patch("unified_sse.HEARTBEAT_INTERVAL", 0.01), + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -461,9 +462,10 @@ def test_max_connection_time(self): def test_cleanup_on_generator_close(self): """Test that client is removed when generator is closed.""" - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -478,9 +480,10 @@ def test_cleanup_on_generator_close(self): def test_empty_snapshot_when_no_repo_path(self): """Test that an empty snapshot is sent when no repo path available.""" - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -491,25 +494,22 @@ def test_empty_snapshot_when_no_repo_path(self): snapshot = next(gen) gen.close() - data_line = [ - l for l in snapshot.split("\n") if l.startswith("data: ") - ][0] + data_line = [l for l in snapshot.split("\n") if l.startswith("data: ")][0] data = json.loads(data_line[6:]) assert data["pipelines"] == [] def test_snapshot_includes_pipeline_metadata(self): """Test that snapshot entries include repo, branch, compact, progress.""" pipelines = [ - create_test_pipeline( - "p-1", repo="owner/repo", branch="egg/fix-bug" - ), + create_test_pipeline("p-1", repo="owner/repo", branch="egg/fix-bug"), ] - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._collect_all_pipelines_safe") as mock_collect, \ - patch("unified_sse.render_compact_status") as mock_compact, \ - patch("unified_sse.render_progress_bar") as mock_progress: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._collect_all_pipelines_safe") as mock_collect, + patch("unified_sse.render_compact_status") as mock_compact, + patch("unified_sse.render_progress_bar") as mock_progress, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -523,9 +523,7 @@ def test_snapshot_includes_pipeline_metadata(self): snapshot = next(gen) gen.close() - data_line = [ - l for l in snapshot.split("\n") if l.startswith("data: ") - ][0] + data_line = [l for l in snapshot.split("\n") if l.startswith("data: ")][0] data = json.loads(data_line[6:]) entry = data["pipelines"][0] assert entry["repo"] == "owner/repo" @@ -533,17 +531,17 @@ def test_snapshot_includes_pipeline_metadata(self): assert entry["compact"] == "[>Implement] -> oPR" assert entry["progress"] == "[####----] 50%" - def test_stream_enriches_events_with_visualization(self): """Test that the stream generator enriches events with compact/progress.""" pipeline = create_test_pipeline("p-1") - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path, \ - patch("unified_sse.get_state_store") as mock_store, \ - patch("unified_sse.render_compact_status") as mock_compact, \ - patch("unified_sse.render_progress_bar") as mock_progress: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + patch("unified_sse.get_state_store") as mock_store, + patch("unified_sse.render_compact_status") as mock_compact, + patch("unified_sse.render_progress_bar") as mock_progress, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -557,17 +555,17 @@ def test_stream_enriches_events_with_visualization(self): mock_progress.return_value = "[##] 50%" # Queue an event payload (as _handle_event would produce) - mock_q.put({ - "event_type": "phase.started", - "pipeline_id": "p-1", - "is_terminal": False, - "timestamp": "2026-01-01T00:00:00Z", - "data": {}, - }) - - gen = create_unified_sse_stream( - repo_path=Path("/tmp/test"), use_ascii=True + mock_q.put( + { + "event_type": "phase.started", + "pipeline_id": "p-1", + "is_terminal": False, + "timestamp": "2026-01-01T00:00:00Z", + "data": {}, + } ) + + gen = create_unified_sse_stream(repo_path=Path("/tmp/test"), use_ascii=True) # Get snapshot next(gen) # Get the enriched event @@ -587,13 +585,14 @@ def test_stream_enriches_full_dag_in_single_load(self): """Test that full_dag enrichment uses a single state store load.""" pipeline = create_test_pipeline("p-1") - with patch("unified_sse.get_unified_sse_manager") as mock_mgr, \ - patch("unified_sse._resolve_repo_path") as mock_path, \ - patch("unified_sse.get_state_store") as mock_store, \ - patch("unified_sse.render_compact_status") as mock_compact, \ - patch("unified_sse.render_progress_bar") as mock_progress, \ - patch("unified_sse.render_pipeline_dag") as mock_dag: - + with ( + patch("unified_sse.get_unified_sse_manager") as mock_mgr, + patch("unified_sse._resolve_repo_path") as mock_path, + patch("unified_sse.get_state_store") as mock_store, + patch("unified_sse.render_compact_status") as mock_compact, + patch("unified_sse.render_progress_bar") as mock_progress, + patch("unified_sse.render_pipeline_dag") as mock_dag, + ): mock_q = Queue() mock_manager = MagicMock() mock_manager.add_client.return_value = mock_q @@ -608,17 +607,17 @@ def test_stream_enriches_full_dag_in_single_load(self): mock_dag.return_value = "PLAN -> IMPLEMENT -> TEST" # Queue an event payload - mock_q.put({ - "event_type": "phase.started", - "pipeline_id": "p-1", - "is_terminal": False, - "timestamp": "2026-01-01T00:00:00Z", - "data": {}, - }) - - gen = create_unified_sse_stream( - repo_path=Path("/tmp/test"), full_dag=True + mock_q.put( + { + "event_type": "phase.started", + "pipeline_id": "p-1", + "is_terminal": False, + "timestamp": "2026-01-01T00:00:00Z", + "data": {}, + } ) + + gen = create_unified_sse_stream(repo_path=Path("/tmp/test"), full_dag=True) # Get snapshot next(gen) # Get the enriched event diff --git a/orchestrator/unified_sse.py b/orchestrator/unified_sse.py index c0dd9cd905..16481fca3f 100644 --- a/orchestrator/unified_sse.py +++ b/orchestrator/unified_sse.py @@ -224,9 +224,7 @@ def create_unified_sse_stream( PipelineStatus.FAILED, PipelineStatus.CANCELLED, } - all_pipelines = [ - p for p in all_pipelines if p.status not in terminal - ] + all_pipelines = [p for p in all_pipelines if p.status not in terminal] snapshot_data: list[dict[str, Any]] = [] for pipeline in all_pipelines: @@ -236,17 +234,11 @@ def create_unified_sse_stream( "current_phase": pipeline.current_phase.value, "repo": pipeline.repo, "branch": pipeline.branch, - "compact": render_compact_status( - pipeline, use_ascii=use_ascii - ), - "progress": render_progress_bar( - pipeline, use_ascii=use_ascii - ), + "compact": render_compact_status(pipeline, use_ascii=use_ascii), + "progress": render_progress_bar(pipeline, use_ascii=use_ascii), } if full_dag: - entry["dag"] = render_pipeline_dag( - pipeline, use_ascii=use_ascii - ) + entry["dag"] = render_pipeline_dag(pipeline, use_ascii=use_ascii) snapshot_data.append(entry) yield format_sse_event( @@ -293,21 +285,13 @@ def create_unified_sse_stream( try: if repo_path: store = get_state_store(repo_path) - pipeline = store.load_pipeline( - payload["pipeline_id"] - ) + pipeline = store.load_pipeline(payload["pipeline_id"]) payload["status"] = pipeline.status.value payload["current_phase"] = pipeline.current_phase.value - payload["compact"] = render_compact_status( - pipeline, use_ascii=use_ascii - ) - payload["progress"] = render_progress_bar( - pipeline, use_ascii=use_ascii - ) + payload["compact"] = render_compact_status(pipeline, use_ascii=use_ascii) + payload["progress"] = render_progress_bar(pipeline, use_ascii=use_ascii) if full_dag: - payload["dag"] = render_pipeline_dag( - pipeline, use_ascii=use_ascii - ) + payload["dag"] = render_pipeline_dag(pipeline, use_ascii=use_ascii) except Exception: logger.debug( "Failed to enrich unified SSE event", @@ -324,9 +308,7 @@ def create_unified_sse_stream( # end the unified stream — we keep watching all pipelines. except Empty: - yield format_sse_comment( - f"heartbeat {datetime.utcnow().isoformat()}Z" - ) + yield format_sse_comment(f"heartbeat {datetime.utcnow().isoformat()}Z") finally: manager.remove_client(q) diff --git a/orchestrator/webhooks.py b/orchestrator/webhooks.py index 1bff26fc59..b3cafd1494 100644 --- a/orchestrator/webhooks.py +++ b/orchestrator/webhooks.py @@ -9,11 +9,10 @@ import hashlib import hmac -import json import sys -from datetime import datetime +from collections.abc import Callable from pathlib import Path -from typing import Any, Callable +from typing import Any from flask import Blueprint, Response, jsonify, request @@ -113,10 +112,12 @@ def github_webhook() -> tuple[Response, int]: signature = request.headers.get("X-Hub-Signature-256") if not event_type: - return jsonify({ - "success": False, - "message": "Missing X-GitHub-Event header", - }), 400 + return jsonify( + { + "success": False, + "message": "Missing X-GitHub-Event header", + } + ), 400 # Verify signature if secret is configured webhook_secret = os.environ.get("GITHUB_WEBHOOK_SECRET") @@ -127,19 +128,23 @@ def github_webhook() -> tuple[Response, int]: event_type=event_type, delivery_id=delivery_id, ) - return jsonify({ - "success": False, - "message": "Invalid signature", - }), 401 + return jsonify( + { + "success": False, + "message": "Invalid signature", + } + ), 401 # Parse payload try: payload = request.get_json() except Exception as e: - return jsonify({ - "success": False, - "message": f"Invalid JSON: {e}", - }), 400 + return jsonify( + { + "success": False, + "message": f"Invalid JSON: {e}", + } + ), 400 logger.info( "GitHub webhook received", @@ -152,11 +157,13 @@ def github_webhook() -> tuple[Response, int]: try: result = handle_github_event(event_type, payload) - return jsonify({ - "success": True, - "message": "Webhook processed", - "data": result, - }), 200 + return jsonify( + { + "success": True, + "message": "Webhook processed", + "data": result, + } + ), 200 except Exception as e: logger.error( @@ -164,10 +171,12 @@ def github_webhook() -> tuple[Response, int]: event_type=event_type, error=str(e), ) - return jsonify({ - "success": False, - "message": f"Processing error: {e}", - }), 500 + return jsonify( + { + "success": False, + "message": f"Processing error: {e}", + } + ), 500 def handle_github_event( @@ -318,20 +327,24 @@ def manual_trigger() -> tuple[Response, int]: """ data = request.get_json() if not data: - return jsonify({ - "success": False, - "message": "Missing request body", - }), 400 + return jsonify( + { + "success": False, + "message": "Missing request body", + } + ), 400 event = data.get("event") issue_number = data.get("issue_number") repo = data.get("repo") if not event: - return jsonify({ - "success": False, - "message": "Missing event type", - }), 400 + return jsonify( + { + "success": False, + "message": "Missing event type", + } + ), 400 logger.info( "Manual trigger received", @@ -361,17 +374,21 @@ def manual_trigger() -> tuple[Response, int]: ) result["handled"] = True - return jsonify({ - "success": True, - "message": "Trigger processed", - "data": result, - }), 200 + return jsonify( + { + "success": True, + "message": "Trigger processed", + "data": result, + } + ), 200 @webhooks_bp.route("/health", methods=["GET"]) def webhook_health() -> tuple[Response, int]: """Webhook endpoint health check.""" - return jsonify({ - "status": "healthy", - "registered_handlers": list(_webhook_handlers.keys()), - }), 200 + return jsonify( + { + "status": "healthy", + "registered_handlers": list(_webhook_handlers.keys()), + } + ), 200 diff --git a/sandbox/entrypoint.py b/sandbox/entrypoint.py index 5d4566716d..0255e32f15 100644 --- a/sandbox/entrypoint.py +++ b/sandbox/entrypoint.py @@ -710,7 +710,9 @@ def setup_sdlc_tokens(config: Config, logger: Logger) -> None: pass project_settings.setdefault("hooks", {}) existing_project_hooks = project_settings["hooks"].get("UserPromptSubmit", []) - if not any(h.get("command") == str(hook_dst) for h in existing_project_hooks if isinstance(h, dict)): + if not any( + h.get("command") == str(hook_dst) for h in existing_project_hooks if isinstance(h, dict) + ): existing_project_hooks.append({"type": "command", "command": str(hook_dst)}) project_settings["hooks"]["UserPromptSubmit"] = existing_project_hooks project_settings_file.write_text(json.dumps(project_settings, indent=2)) @@ -721,9 +723,11 @@ def setup_sdlc_tokens(config: Config, logger: Logger) -> None: claude_md = config.user_home / "CLAUDE.md" if claude_md.exists(): with open(claude_md, "a") as f: - f.write(f"\n\n---\n\n## Auto-Start SDLC\n\n" - f"IMMEDIATELY run `/sdlc {sdlc_issue}` as your first action. " - f"Do not wait for user input.\n") + f.write( + f"\n\n---\n\n## Auto-Start SDLC\n\n" + f"IMMEDIATELY run `/sdlc {sdlc_issue}` as your first action. " + f"Do not wait for user input.\n" + ) logger.success(f"Auto-start instruction added to CLAUDE.md for issue {sdlc_issue}") # Start settings watchdog (background thread, runs as root) @@ -750,9 +754,7 @@ def watchdog() -> None: # Check if our hook is present hook_present = any( - h.get("command") == hook_path - for h in user_prompt_hooks - if isinstance(h, dict) + h.get("command") == hook_path for h in user_prompt_hooks if isinstance(h, dict) ) if not hook_present: diff --git a/shared/egg_contracts/checkpoint_cli.py b/shared/egg_contracts/checkpoint_cli.py index e18cdde678..a8f0dad733 100644 --- a/shared/egg_contracts/checkpoint_cli.py +++ b/shared/egg_contracts/checkpoint_cli.py @@ -3,17 +3,20 @@ Checkpoint CLI for browsing and querying agent checkpoints. This CLI provides commands to list, show, and browse checkpoints that capture -agent session context. Each commit has exactly one checkpoint, enabling precise -traceability between code changes and agent sessions. +agent session context. Supports both commit-triggered and session-end checkpoints +with rich multi-dimensional querying. -Checkpoints are stored in the egg/checkpoints/v1 branch and are captured -per-commit during git push operations. Transcript data is extracted from the -API proxy buffer, providing stable and format-independent capture. +Checkpoints are stored in the egg/checkpoints/v2 branch and are captured +per-commit during git push operations and at session end. Transcript data is +extracted from the API proxy buffer, providing stable and format-independent +capture. Commands: egg-checkpoint list [--branch ] [--issue ] [--limit ] + [--trigger ] [--status ] [--agent-type ] + [--session ] [--pr ] [--phase ] List checkpoints with metadata - egg-checkpoint show Display full checkpoint details + egg-checkpoint show Display full checkpoint details egg-checkpoint browse --issue Filter checkpoints by issue """ @@ -28,16 +31,20 @@ from typing import Any from .checkpoint_loader import ( - CheckpointLoadError, - list_checkpoints, - load_checkpoint, - load_checkpoint_by_commit, - load_checkpoint_index, + list_checkpoints_v2, + load_checkpoint_by_commit_v2, + load_checkpoint_by_id_v2, +) +from .checkpoints import ( + AgentType, + CheckpointSummaryV2, + CheckpointV2, + SessionStatus, + TriggerType, ) -from .checkpoints import Checkpoint, CheckpointIndex # Checkpoint branch name -CHECKPOINT_BRANCH = "egg/checkpoints/v1" +CHECKPOINT_BRANCH = "egg/checkpoints/v2" def get_repo_path() -> str: @@ -45,7 +52,9 @@ def get_repo_path() -> str: return os.environ.get("EGG_REPO_PATH", str(Path.cwd())) -def run_git(args: list[str], cwd: str | None = None, check: bool = True) -> subprocess.CompletedProcess[str]: +def run_git( + args: list[str], cwd: str | None = None, check: bool = True +) -> subprocess.CompletedProcess[str]: """Run a git command.""" cmd = ["git"] + args result = subprocess.run( @@ -93,6 +102,7 @@ def checkout_checkpoint_branch(repo_path: str) -> Path | None: # Cleanup on failure if temp_path.exists(): import shutil + shutil.rmtree(temp_path, ignore_errors=True) return None @@ -107,6 +117,7 @@ def cleanup_worktree(repo_path: str, worktree_path: Path) -> None: # Also try to remove the directory if worktree remove failed if worktree_path.exists(): import shutil + shutil.rmtree(worktree_path, ignore_errors=True) @@ -131,42 +142,74 @@ def format_tokens(n: int) -> str: return str(n) -def print_checkpoint_summary(checkpoint: Checkpoint | dict[str, Any]) -> None: +def _format_trigger(trigger_type: TriggerType | str) -> str: + """Format trigger type for display.""" + value = trigger_type.value if isinstance(trigger_type, TriggerType) else trigger_type + return {"commit": "commit", "session_end": "session-end"}.get(value, value) + + +def _format_status(status: SessionStatus | str | None) -> str: + """Format session status for display.""" + if status is None: + return "" + return status.value if isinstance(status, SessionStatus) else status + + +def print_checkpoint_summary( + checkpoint: CheckpointSummaryV2 | CheckpointV2 | dict[str, Any], +) -> None: """Print a one-line summary of a checkpoint.""" - if isinstance(checkpoint, Checkpoint): + if isinstance(checkpoint, (CheckpointSummaryV2, CheckpointV2)): data = checkpoint.model_dump() else: data = checkpoint cp_id = data.get("id", "unknown") - commit = data.get("commit_sha", "unknown")[:7] + commit = data.get("commit_sha") created = format_timestamp(data.get("created_at")) branch = data.get("branch", "") issue = data.get("issue_number") + pr = data.get("pr_number") phase = data.get("pipeline_phase", "") - - session = data.get("session", {}) - role = session.get("agent_role", "") - - transcript = data.get("transcript", {}) - msg_count = transcript.get("message_count", 0) if transcript else 0 - - tool_calls = data.get("tool_calls", []) - tool_count = len(tool_calls) - - token_usage = data.get("token_usage", {}) - tokens = token_usage.get("total_tokens", 0) if token_usage else 0 + trigger = data.get("trigger_type", "") + status = data.get("session_status", "") + agent = data.get("agent_type", "") + + # Extract metrics depending on model type + if isinstance(checkpoint, CheckpointV2): + msg_count = checkpoint.transcript.message_count if checkpoint.transcript else 0 + tool_count = len(checkpoint.tool_calls) + tokens = checkpoint.token_usage.total_tokens if checkpoint.token_usage else 0 + elif isinstance(checkpoint, CheckpointSummaryV2): + msg_count = checkpoint.message_count + tool_count = checkpoint.tool_call_count + tokens = checkpoint.total_tokens + else: + transcript = data.get("transcript", {}) + msg_count = data.get( + "message_count", transcript.get("message_count", 0) if transcript else 0 + ) + tool_count = data.get("tool_call_count", len(data.get("tool_calls", []))) + token_usage = data.get("token_usage", {}) + tokens = data.get("total_tokens", token_usage.get("total_tokens", 0) if token_usage else 0) # Build summary line - parts = [f"{cp_id}", f"commit:{commit}"] + parts = [f"{cp_id}"] + if commit: + parts.append(f"commit:{commit[:7]}") + parts.append(f"trigger:{_format_trigger(trigger)}") + if status: + parts.append(f"status:{_format_status(status)}") if branch: parts.append(f"branch:{branch}") if issue: parts.append(f"issue:#{issue}") + if pr: + parts.append(f"pr:#{pr}") if phase: parts.append(f"phase:{phase}") - if role: - parts.append(f"role:{role}") + if agent and agent != "unknown": + parts.append(f"agent:{agent}") parts.append(f"msgs:{msg_count}") parts.append(f"tools:{tool_count}") parts.append(f"tokens:{format_tokens(tokens)}") @@ -175,37 +218,53 @@ def print_checkpoint_summary(checkpoint: Checkpoint | dict[str, Any]) -> None: print(" " + " | ".join(parts)) -def print_checkpoint_details(checkpoint: Checkpoint | dict[str, Any]) -> None: +def print_checkpoint_details(checkpoint: CheckpointV2 | dict[str, Any]) -> None: """Print detailed checkpoint information.""" - if isinstance(checkpoint, Checkpoint): + if isinstance(checkpoint, CheckpointV2): data = checkpoint.model_dump() else: data = checkpoint print(f"Checkpoint: {data.get('id')}") - print(f" Commit: {data.get('commit_sha')}") + print(f" Trigger: {_format_trigger(data.get('trigger_type', ''))}") + if data.get("session_status"): + print(f" Session Status: {_format_status(data.get('session_status'))}") + if data.get("commit_sha"): + print(f" Commit: {data.get('commit_sha')}") + if data.get("push_sha"): + print(f" Push SHA: {data.get('push_sha')}") print(f" Branch: {data.get('branch', 'N/A')}") print(f" Created: {format_timestamp(data.get('created_at'))}") if data.get("issue_number"): print(f" Issue: #{data.get('issue_number')}") + if data.get("pr_number"): + print(f" PR: #{data.get('pr_number')}") if data.get("pipeline_phase"): print(f" Phase: {data.get('pipeline_phase')}") + agent = data.get("agent_type", "") + if agent and agent != "unknown": + print(f" Agent Type: {agent}") + # Session metadata print() print("Session:") + print(f" Session ID: {data.get('session_id', 'N/A')}") session = data.get("session", {}) - print(f" Session ID: {session.get('session_id', 'N/A')}") if session.get("container_id"): print(f" Container: {session.get('container_id')}") if session.get("agent_role"): print(f" Role: {session.get('agent_role')}") if session.get("model"): print(f" Model: {session.get('model')}") - print(f" Started: {format_timestamp(session.get('started_at'))}") - if session.get("ended_at"): - print(f" Ended: {format_timestamp(session.get('ended_at'))}") + print( + f" Started: {format_timestamp(data.get('session_started_at') or session.get('started_at'))}" + ) + if data.get("session_ended_at") or session.get("ended_at"): + print( + f" Ended: {format_timestamp(data.get('session_ended_at') or session.get('ended_at'))}" + ) if session.get("duration_seconds"): duration = session.get("duration_seconds") if duration >= 3600: @@ -275,33 +334,34 @@ def cmd_list(args: argparse.Namespace) -> int: try: checkpoints_dir = worktree_path / "checkpoints" + index_path = worktree_path / "index.json" - # Load checkpoints - checkpoints = list_checkpoints( + summaries = list_checkpoints_v2( checkpoints_dir, + index_path, issue_number=args.issue, + pr_number=getattr(args, "pr", None), branch=args.branch, + session_id=getattr(args, "session", None), + trigger_type=getattr(args, "trigger", None), + session_status=getattr(args, "status", None), + agent_type=getattr(args, "agent_type", None), + pipeline_phase=getattr(args, "phase", None), limit=args.limit, ) - if not checkpoints: - filters = [] - if args.issue: - filters.append(f"issue #{args.issue}") - if args.branch: - filters.append(f"branch '{args.branch}'") - filter_str = f" matching {', '.join(filters)}" if filters else "" - print(f"No checkpoints found{filter_str}") + if not summaries: + print("No checkpoints found matching filters") return 0 if args.json: - output = [cp.model_dump(mode="json") for cp in checkpoints] + output = [s.model_dump(mode="json") for s in summaries] print(json.dumps(output, indent=2)) else: - print(f"Checkpoints ({len(checkpoints)} found):") + print(f"Checkpoints ({len(summaries)} found):") print() - for cp in checkpoints: - print_checkpoint_summary(cp) + for s in summaries: + print_checkpoint_summary(s) return 0 @@ -310,9 +370,9 @@ def cmd_list(args: argparse.Namespace) -> int: def cmd_show(args: argparse.Namespace) -> int: - """Display full checkpoint details for a commit.""" + """Display full checkpoint details by checkpoint ID or commit SHA.""" repo_path = args.repo_path or get_repo_path() - commit_sha = args.commit + identifier = args.identifier # Checkout checkpoint branch worktree_path = checkout_checkpoint_branch(repo_path) @@ -324,15 +384,17 @@ def cmd_show(args: argparse.Namespace) -> int: checkpoints_dir = worktree_path / "checkpoints" index_path = worktree_path / "index.json" - # Load checkpoint by commit - checkpoint = load_checkpoint_by_commit( - commit_sha, - checkpoints_dir, - index_path, - ) + checkpoint: CheckpointV2 | None = None + + # Try as checkpoint ID first (ckpt-... prefix) + if identifier.startswith("ckpt-"): + checkpoint = load_checkpoint_by_id_v2(identifier, checkpoints_dir) + else: + # Try as commit SHA + checkpoint = load_checkpoint_by_commit_v2(identifier, checkpoints_dir, index_path) if not checkpoint: - print(f"No checkpoint found for commit {commit_sha}") + print(f"No checkpoint found for '{identifier}'") return 1 if args.json: @@ -358,39 +420,45 @@ def cmd_browse(args: argparse.Namespace) -> int: try: checkpoints_dir = worktree_path / "checkpoints" + index_path = worktree_path / "index.json" - # Load checkpoints for issue - checkpoints = list_checkpoints( + summaries = list_checkpoints_v2( checkpoints_dir, + index_path, issue_number=args.issue, limit=args.limit, ) - if not checkpoints: + if not summaries: print(f"No checkpoints found for issue #{args.issue}") return 0 if args.json: - output = [cp.model_dump(mode="json") for cp in checkpoints] + output = [s.model_dump(mode="json") for s in summaries] print(json.dumps(output, indent=2)) else: - print(f"Checkpoints for Issue #{args.issue} ({len(checkpoints)} found):") + print(f"Checkpoints for Issue #{args.issue} ({len(summaries)} found):") print() # Group by session - sessions: dict[str, list[Checkpoint]] = {} - for cp in checkpoints: - session_id = cp.session.session_id - if session_id not in sessions: - sessions[session_id] = [] - sessions[session_id].append(cp) - - for session_id, session_checkpoints in sessions.items(): - first = session_checkpoints[0] - role = first.session.agent_role or "unknown" - print(f"Session: {session_id[:12]}... (role: {role})") - for cp in session_checkpoints: - print_checkpoint_summary(cp) + sessions: dict[str, list[CheckpointSummaryV2]] = {} + for s in summaries: + sid = s.session_id + if sid not in sessions: + sessions[sid] = [] + sessions[sid].append(s) + + for sid, session_summaries in sessions.items(): + first = session_summaries[0] + agent = ( + first.agent_type.value if first.agent_type != AgentType.UNKNOWN else "unknown" + ) + triggers = {_format_trigger(s.trigger_type) for s in session_summaries} + print( + f"Session: {sid[:12]}... (agent: {agent}, triggers: {', '.join(sorted(triggers))})" + ) + for s in session_summaries: + print_checkpoint_summary(s) print() return 0 @@ -416,13 +484,35 @@ def create_parser() -> argparse.ArgumentParser: list_parser = subparsers.add_parser("list", help="List checkpoints with metadata") list_parser.add_argument("--branch", help="Filter by branch name") list_parser.add_argument("--issue", type=int, help="Filter by issue number") + list_parser.add_argument("--pr", type=int, help="Filter by PR number") + list_parser.add_argument("--session", help="Filter by session ID") + list_parser.add_argument( + "--trigger", + choices=[t.value for t in TriggerType], + help="Filter by trigger type", + ) + list_parser.add_argument( + "--status", + choices=[s.value for s in SessionStatus], + help="Filter by session status", + ) + list_parser.add_argument( + "--agent-type", + choices=[a.value for a in AgentType], + help="Filter by agent type", + ) + list_parser.add_argument( + "--phase", + choices=["refine", "plan", "implement", "pr"], + help="Filter by pipeline phase", + ) list_parser.add_argument("--limit", type=int, default=50, help="Maximum checkpoints to show") list_parser.add_argument("--json", action="store_true", help="Output as JSON") list_parser.set_defaults(func=cmd_list) # show command show_parser = subparsers.add_parser("show", help="Display full checkpoint details") - show_parser.add_argument("commit", help="Commit SHA to show checkpoint for") + show_parser.add_argument("identifier", help="Checkpoint ID (ckpt-...) or commit SHA") show_parser.add_argument("--json", action="store_true", help="Output as JSON") show_parser.set_defaults(func=cmd_show) diff --git a/shared/egg_contracts/checkpoint_loader.py b/shared/egg_contracts/checkpoint_loader.py index bb38069181..498d5b82de 100644 --- a/shared/egg_contracts/checkpoint_loader.py +++ b/shared/egg_contracts/checkpoint_loader.py @@ -8,12 +8,18 @@ import hashlib import json import os -import secrets import tempfile from datetime import UTC, datetime from pathlib import Path -from .checkpoints import Checkpoint, CheckpointIndex, CheckpointSummary +from .checkpoints import ( + AgentType, + CheckpointIndexV2, + CheckpointSummaryV2, + CheckpointV2, + SessionStatus, + TriggerType, +) class CheckpointLoadError(Exception): @@ -28,18 +34,6 @@ class CheckpointSaveError(Exception): pass -def generate_checkpoint_id() -> str: - """ - Generate a unique checkpoint ID. - - Returns: - A checkpoint ID in the format ckpt-{12 hex chars} - """ - random_bytes = secrets.token_bytes(6) - hex_str = random_bytes.hex() - return f"ckpt-{hex_str}" - - def generate_checkpoint_id_from_commit( commit_sha: str, session_id: str, @@ -101,95 +95,86 @@ def get_checkpoint_path(base_dir: Path, checkpoint_id: str) -> Path: return base_dir / prefix / get_checkpoint_filename(checkpoint_id) -def load_checkpoint(path: Path) -> Checkpoint: +# ============================================================================== +# Checkpoint v2 Loader Functions +# +# V2 checkpoints support session-end triggers (no commit_sha required) and +# multi-dimensional indexing. Stored on the egg/checkpoints/v2 branch. +# ============================================================================== + + +def generate_checkpoint_id_v2( + session_id: str, + timestamp: datetime | None = None, +) -> str: """ - Load a checkpoint from a JSON file. + Generate a deterministic checkpoint ID for session-end checkpoints. + + Unlike commit-based IDs, these are derived from session_id + timestamp + since no commit_sha is available. Args: - path: Path to the checkpoint JSON file + session_id: The session/container ID + timestamp: Optional timestamp for uniqueness (defaults to now) Returns: - The loaded Checkpoint - - Raises: - CheckpointLoadError: If the file cannot be loaded or parsed + A checkpoint ID in the format ckpt-{16 hex chars} """ - try: - with open(path) as f: - data = json.load(f) - return Checkpoint.model_validate(data) - except FileNotFoundError as e: - msg = f"Checkpoint file not found: {path}" - raise CheckpointLoadError(msg) from e - except json.JSONDecodeError as e: - msg = f"Invalid JSON in checkpoint file {path}: {e}" - raise CheckpointLoadError(msg) from e - except Exception as e: - msg = f"Error loading checkpoint from {path}: {e}" - raise CheckpointLoadError(msg) from e + if timestamp is None: + timestamp = datetime.now(UTC) + content = f"session:{session_id}:{timestamp.isoformat()}" + hash_bytes = hashlib.sha256(content.encode()).digest()[:8] + hex_str = hash_bytes.hex() + return f"ckpt-{hex_str}" -def save_checkpoint(checkpoint: Checkpoint, path: Path) -> None: +def save_checkpoint_v2(checkpoint: CheckpointV2, path: Path) -> None: """ - Save a checkpoint to a JSON file atomically. + Save a v2 checkpoint to a JSON file atomically. - Uses the temp file + rename pattern to ensure atomic writes: - 1. Write to a temporary file in the same directory - 2. Sync to disk - 3. Atomically rename to the target path - - This prevents corruption if the process is interrupted during write. + Uses the same temp file + rename pattern as v1 for crash safety. Args: - checkpoint: The checkpoint to save + checkpoint: The v2 checkpoint to save path: Path to save the checkpoint to Raises: CheckpointSaveError: If the checkpoint cannot be saved """ try: - # Ensure parent directory exists path.parent.mkdir(parents=True, exist_ok=True) - # Serialize checkpoint to JSON data = checkpoint.model_dump(mode="json") json_str = json.dumps(data, indent=2, sort_keys=True) - # Write to temp file in same directory (for atomic rename) - fd, temp_path = tempfile.mkstemp( - suffix=".tmp", prefix=".checkpoint_", dir=path.parent - ) + fd, temp_path = tempfile.mkstemp(suffix=".tmp", prefix=".checkpoint_", dir=path.parent) try: with os.fdopen(fd, "w") as f: f.write(json_str) f.flush() os.fsync(f.fileno()) - # Set restrictive permissions os.chmod(temp_path, 0o644) - - # Atomic rename os.rename(temp_path, path) except Exception: - # Clean up temp file on error if os.path.exists(temp_path): os.unlink(temp_path) raise except Exception as e: - msg = f"Error saving checkpoint to {path}: {e}" + msg = f"Error saving v2 checkpoint to {path}: {e}" raise CheckpointSaveError(msg) from e -def load_checkpoint_index(path: Path) -> CheckpointIndex: +def load_checkpoint_v2(path: Path) -> CheckpointV2: """ - Load the checkpoint index from a JSON file. + Load a v2 checkpoint from a JSON file. Args: - path: Path to the index JSON file + path: Path to the checkpoint JSON file Returns: - The loaded CheckpointIndex + The loaded CheckpointV2 Raises: CheckpointLoadError: If the file cannot be loaded or parsed @@ -197,38 +182,61 @@ def load_checkpoint_index(path: Path) -> CheckpointIndex: try: with open(path) as f: data = json.load(f) - return CheckpointIndex.model_validate(data) + return CheckpointV2.model_validate(data) + except FileNotFoundError as e: + msg = f"V2 checkpoint file not found: {path}" + raise CheckpointLoadError(msg) from e + except json.JSONDecodeError as e: + msg = f"Invalid JSON in v2 checkpoint file {path}: {e}" + raise CheckpointLoadError(msg) from e + except Exception as e: + msg = f"Error loading v2 checkpoint from {path}: {e}" + raise CheckpointLoadError(msg) from e + + +def load_checkpoint_index_v2(path: Path) -> CheckpointIndexV2: + """ + Load the v2 checkpoint index from a JSON file. + + Returns an empty index if the file doesn't exist. + + Args: + path: Path to the index JSON file + + Returns: + The loaded CheckpointIndexV2 + """ + try: + with open(path) as f: + data = json.load(f) + return CheckpointIndexV2.model_validate(data) except FileNotFoundError: - # Return empty index if file doesn't exist - return CheckpointIndex(last_updated=datetime.now(UTC), checkpoints=[]) + return CheckpointIndexV2(last_updated=datetime.now(UTC)) except json.JSONDecodeError as e: - msg = f"Invalid JSON in index file {path}: {e}" + msg = f"Invalid JSON in v2 index file {path}: {e}" raise CheckpointLoadError(msg) from e except Exception as e: - msg = f"Error loading checkpoint index from {path}: {e}" + msg = f"Error loading v2 checkpoint index from {path}: {e}" raise CheckpointLoadError(msg) from e -def save_checkpoint_index(index: CheckpointIndex, path: Path) -> None: +def save_checkpoint_index_v2(index: CheckpointIndexV2, path: Path) -> None: """ - Save the checkpoint index to a JSON file atomically. + Save the v2 checkpoint index to a JSON file atomically. Args: - index: The checkpoint index to save + index: The v2 checkpoint index to save path: Path to save the index to Raises: CheckpointSaveError: If the index cannot be saved """ try: - # Ensure parent directory exists path.parent.mkdir(parents=True, exist_ok=True) - # Serialize index to JSON data = index.model_dump(mode="json") json_str = json.dumps(data, indent=2, sort_keys=True) - # Write to temp file in same directory (for atomic rename) fd, temp_path = tempfile.mkstemp(suffix=".tmp", prefix=".index_", dir=path.parent) try: with os.fdopen(fd, "w") as f: @@ -236,165 +244,227 @@ def save_checkpoint_index(index: CheckpointIndex, path: Path) -> None: f.flush() os.fsync(f.fileno()) - # Set restrictive permissions os.chmod(temp_path, 0o644) - - # Atomic rename os.rename(temp_path, path) except Exception: - # Clean up temp file on error if os.path.exists(temp_path): os.unlink(temp_path) raise except Exception as e: - msg = f"Error saving checkpoint index to {path}: {e}" + msg = f"Error saving v2 checkpoint index to {path}: {e}" raise CheckpointSaveError(msg) from e -def add_checkpoint_to_index( - checkpoint: Checkpoint, +def add_checkpoint_to_index_v2( + checkpoint: CheckpointV2, index_path: Path, -) -> CheckpointIndex: +) -> CheckpointIndexV2: """ - Add a checkpoint to the index and save it. + Add a v2 checkpoint to the index with multi-dimensional index updates. - This is a convenience function that: - 1. Loads the existing index (or creates a new one) - 2. Adds the checkpoint summary - 3. Saves the updated index atomically + Updates the primary checkpoint list and all secondary indices. The final + file write uses atomic rename, but the read-modify-write cycle is not + locked — concurrent callers must coordinate externally (e.g., via git + push non-fast-forward rejection) to avoid lost updates. + Deduplicates by checkpoint ID. Args: - checkpoint: The checkpoint to add + checkpoint: The v2 checkpoint to add index_path: Path to the index file Returns: - The updated CheckpointIndex + The updated CheckpointIndexV2 Raises: CheckpointSaveError: If the index cannot be saved """ - # Load existing index or create new one - index = load_checkpoint_index(index_path) + index = load_checkpoint_index_v2(index_path) - # Check if checkpoint already exists in index - existing = index.get_by_commit(checkpoint.commit_sha) - if existing and existing.id == checkpoint.id: - # Already in index with same ID, no update needed + # Deduplicate by checkpoint ID + existing_ids = {cp.id for cp in index.checkpoints} + if checkpoint.id in existing_ids: return index - # Remove any existing entry for this commit (in case of re-checkpoint) - index.checkpoints = [cp for cp in index.checkpoints if cp.commit_sha != checkpoint.commit_sha] - - # Add new checkpoint summary - summary = CheckpointSummary.from_checkpoint(checkpoint) + # Create summary and add to primary list + summary = CheckpointSummaryV2.from_checkpoint(checkpoint) index.checkpoints.append(summary) - # Update timestamp + # Update secondary indices + _append_to_index(index.by_session, checkpoint.session_id, checkpoint.id) + _append_to_index(index.by_trigger, checkpoint.trigger_type.value, checkpoint.id) + + if checkpoint.issue_number is not None: + _append_to_index(index.by_issue, str(checkpoint.issue_number), checkpoint.id) + + if checkpoint.pr_number is not None: + _append_to_index(index.by_pr, str(checkpoint.pr_number), checkpoint.id) + + if checkpoint.commit_sha is not None: + index.by_commit[checkpoint.commit_sha] = checkpoint.id + + _append_to_index(index.by_agent_type, checkpoint.agent_type.value, checkpoint.id) + + if checkpoint.pipeline_phase is not None: + _append_to_index(index.by_phase, checkpoint.pipeline_phase, checkpoint.id) + + if checkpoint.session_status is not None: + _append_to_index(index.by_status, checkpoint.session_status.value, checkpoint.id) + index.last_updated = datetime.now(UTC) - # Save atomically - save_checkpoint_index(index, index_path) + save_checkpoint_index_v2(index, index_path) return index -def load_checkpoint_by_commit( - commit_sha: str, +def _append_to_index(index_dict: dict[str, list[str]], key: str, value: str) -> None: + """Append value to index dict list, deduplicating.""" + if key not in index_dict: + index_dict[key] = [] + if value not in index_dict[key]: + index_dict[key].append(value) + + +def load_checkpoint_by_id_v2( + checkpoint_id: str, checkpoints_dir: Path, - index_path: Path | None = None, -) -> Checkpoint | None: +) -> CheckpointV2 | None: """ - Load a checkpoint by commit SHA. + Load a v2 checkpoint by its ID. + + Args: + checkpoint_id: The checkpoint ID (e.g., ckpt-abc123def456) + checkpoints_dir: Directory containing checkpoint files + + Returns: + The CheckpointV2 if found, None otherwise + """ + checkpoint_path = get_checkpoint_path(checkpoints_dir, checkpoint_id) + if not checkpoint_path.exists(): + return None + try: + return load_checkpoint_v2(checkpoint_path) + except CheckpointLoadError: + return None - If an index path is provided, uses the index for fast lookup. - Otherwise, scans the checkpoint directory. + +def load_checkpoint_by_commit_v2( + commit_sha: str, + checkpoints_dir: Path, + index_path: Path, +) -> CheckpointV2 | None: + """ + Load a v2 checkpoint by commit SHA using the v2 index. Args: commit_sha: The commit SHA to find checkpoints_dir: Directory containing checkpoint files - index_path: Optional path to the checkpoint index + index_path: Path to the v2 index file Returns: - The Checkpoint if found, None otherwise + The CheckpointV2 if found, None otherwise """ - # Try using index for fast lookup - if index_path and index_path.exists(): - try: - index = load_checkpoint_index(index_path) - summary = index.get_by_commit(commit_sha) - if summary: - checkpoint_path = get_checkpoint_path(checkpoints_dir, summary.id) - if checkpoint_path.exists(): - return load_checkpoint(checkpoint_path) - except CheckpointLoadError: - pass # Fall through to scan - - # Fallback: scan checkpoint directory - if not checkpoints_dir.exists(): + try: + index = load_checkpoint_index_v2(index_path) + except CheckpointLoadError: return None - for subdir in checkpoints_dir.iterdir(): - if not subdir.is_dir(): - continue - for checkpoint_file in subdir.glob("ckpt-*.json"): - try: - checkpoint = load_checkpoint(checkpoint_file) - if ( - checkpoint.commit_sha == commit_sha - or commit_sha.startswith(checkpoint.commit_sha[:7]) - ): - return checkpoint - except CheckpointLoadError: - continue + checkpoint_id = index.get_by_commit(commit_sha) + if not checkpoint_id: + return None - return None + return load_checkpoint_by_id_v2(checkpoint_id, checkpoints_dir) -def list_checkpoints( +def list_checkpoints_v2( checkpoints_dir: Path, + index_path: Path, issue_number: int | None = None, + pr_number: int | None = None, branch: str | None = None, + session_id: str | None = None, + trigger_type: str | None = None, + session_status: str | None = None, + agent_type: str | None = None, + pipeline_phase: str | None = None, limit: int | None = None, -) -> list[Checkpoint]: +) -> list[CheckpointSummaryV2]: """ - List checkpoints, optionally filtered by issue or branch. + List v2 checkpoint summaries using the index, with multi-dimensional filtering. + + Uses the v2 index for fast lookups. Filters are intersected (AND logic). Args: - checkpoints_dir: Directory containing checkpoint files - issue_number: Optional issue number to filter by - branch: Optional branch to filter by - limit: Optional maximum number of checkpoints to return + checkpoints_dir: Directory containing checkpoint files (unused, kept for API symmetry) + index_path: Path to the v2 index file + issue_number: Filter by issue number + pr_number: Filter by PR number + branch: Filter by branch name + session_id: Filter by session ID + trigger_type: Filter by trigger type value + session_status: Filter by session status value + agent_type: Filter by agent type value + pipeline_phase: Filter by pipeline phase + limit: Maximum number of results Returns: - List of Checkpoint objects, sorted by created_at descending + List of CheckpointSummaryV2, sorted by created_at descending """ - if not checkpoints_dir.exists(): + try: + index = load_checkpoint_index_v2(index_path) + except CheckpointLoadError: return [] - checkpoints = [] - for subdir in checkpoints_dir.iterdir(): - if not subdir.is_dir(): - continue - for checkpoint_file in subdir.glob("ckpt-*.json"): - try: - checkpoint = load_checkpoint(checkpoint_file) + if not index.checkpoints: + return [] + + # Build set of matching checkpoint IDs using index lookups + # Start with None (meaning "all") and intersect with each filter + matching_ids: set[str] | None = None + + def _intersect(ids: list[str]) -> None: + nonlocal matching_ids + id_set = set(ids) + if matching_ids is None: + matching_ids = id_set + else: + matching_ids &= id_set + + if issue_number is not None: + _intersect(index.get_by_issue(issue_number)) + + if pr_number is not None: + _intersect(index.get_by_pr(pr_number)) - # Apply filters - if issue_number is not None and checkpoint.issue_number != issue_number: - continue - if branch is not None and checkpoint.branch != branch: - continue + if session_id is not None: + _intersect(index.get_by_session(session_id)) - checkpoints.append(checkpoint) - except CheckpointLoadError: - continue + if trigger_type is not None: + _intersect(index.get_by_trigger(TriggerType(trigger_type))) + + if session_status is not None: + _intersect(index.get_by_status(SessionStatus(session_status))) + + if agent_type is not None: + _intersect(index.get_by_agent_type(AgentType(agent_type))) + + if pipeline_phase is not None: + _intersect(index.get_by_phase(pipeline_phase)) + + # Filter summaries + results = [] + for summary in index.checkpoints: + if matching_ids is not None and summary.id not in matching_ids: + continue + if branch is not None and summary.branch != branch: + continue + results.append(summary) - # Sort by created_at descending (most recent first) - checkpoints.sort(key=lambda cp: cp.created_at, reverse=True) + # Sort by created_at descending + results.sort(key=lambda cp: cp.created_at, reverse=True) - # Apply limit if limit is not None: - checkpoints = checkpoints[:limit] + results = results[:limit] - return checkpoints + return results diff --git a/shared/egg_contracts/checkpoints.py b/shared/egg_contracts/checkpoints.py index 7a92d102d8..1273ff8493 100644 --- a/shared/egg_contracts/checkpoints.py +++ b/shared/egg_contracts/checkpoints.py @@ -128,53 +128,105 @@ def calculate_total(self) -> int: return self.input_tokens + self.output_tokens -class Checkpoint(BaseModel): +# ============================================================================== +# Checkpoint v2 Models +# +# V2 checkpoints capture every agent session (not just push events) and support +# rich multi-dimensional querying. Stored on the egg/checkpoints/v2 branch. +# ============================================================================== + + +class TriggerType(StrEnum): + """What triggered checkpoint creation.""" + + COMMIT = "commit" + SESSION_END = "session_end" + + +class SessionStatus(StrEnum): + """Terminal state of the session.""" + + COMPLETED = "completed" + EXPIRED = "expired" + FAILED = "failed" + + +class AgentType(StrEnum): + """Agent role/type for classification.""" + + CODER = "coder" + TESTER = "tester" + DOCUMENTER = "documenter" + INTEGRATOR = "integrator" + REVIEWER = "reviewer" + UNKNOWN = "unknown" + + +class CheckpointV2(BaseModel): """ - An agent checkpoint capturing session context. + V2 checkpoint with rich metadata for querying. - Checkpoints are created when agents push commits and capture the full - reasoning context: transcript, tool calls, files touched, and token usage. - They are stored in the egg/checkpoints/v1 branch and linked to commits. + Unlike v1, commit_sha is optional (session-end checkpoints may not have + commits) and trigger_type/session_id are required at top level for + direct indexing. """ schemaVersion: str = Field( # noqa: N815 - default="1.0", pattern=r"^[0-9]+\.[0-9]+$", description="Schema version" + default="2.0", pattern=r"^[0-9]+\.[0-9]+$", description="Schema version" ) id: str = Field( ..., pattern=r"^ckpt-[a-f0-9]{8,16}$", description="Unique checkpoint identifier" ) - commit_sha: str = Field( - ..., - pattern=r"^[a-f0-9]{7,40}$", - description="Git commit SHA this checkpoint is associated with", - ) - session: SessionMetadata = Field(..., description="Session metadata") - transcript: Transcript | None = Field( - default=None, description="The conversation transcript" + + # Trigger context + trigger_type: TriggerType = Field(..., description="What created this checkpoint") + session_status: SessionStatus | None = Field( + default=None, description="Terminal session state (only for SESSION_END triggers)" ) - files_touched: list[FileOperation] = Field( - default_factory=list, description="Files that were read, created, or edited" + + # Git context (optional for session-end checkpoints) + commit_sha: str | None = Field( + default=None, + pattern=r"^[a-f0-9]{7,40}$", + description="Git commit SHA (optional in v2)", ) - tool_calls: list[ToolCall] = Field( - default_factory=list, description="Tool invocations made during the session" + push_sha: str | None = Field( + default=None, + pattern=r"^[a-f0-9]{7,40}$", + description="The tip commit SHA of the push", ) - token_usage: TokenUsage | None = Field(default=None, description="Token usage for the session") + branch: str | None = Field(default=None, description="Git branch") + + # Workflow correlation + session_id: str = Field(..., description="Session/container ID for direct indexing") issue_number: int | None = Field( default=None, ge=1, description="GitHub issue number if associated" ) + pr_number: int | None = Field(default=None, ge=1, description="GitHub PR number if associated") + + # Agent classification + agent_type: AgentType = Field( + default=AgentType.UNKNOWN, description="Agent role classification" + ) pipeline_phase: str | None = Field( default=None, description="SDLC pipeline phase when checkpoint was created" ) - branch: str | None = Field(default=None, description="Git branch where the commit was made") - created_at: datetime = Field(..., description="When checkpoint was created") - push_sha: str | None = Field( - default=None, - pattern=r"^[a-f0-9]{7,40}$", - description="The tip commit SHA of the push", + + # Session details + session: SessionMetadata = Field(..., description="Session metadata") + transcript: Transcript | None = Field(default=None, description="The conversation transcript") + files_touched: list[FileOperation] = Field( + default_factory=list, description="Files that were read, created, or edited" ) - pr_number: int | None = Field( - default=None, ge=1, description="GitHub PR number if associated" + tool_calls: list[ToolCall] = Field( + default_factory=list, description="Tool invocations made during the session" ) + token_usage: TokenUsage | None = Field(default=None, description="Token usage for the session") + + # Timestamps + created_at: datetime = Field(..., description="When checkpoint was created") + session_started_at: datetime = Field(..., description="When session started") + session_ended_at: datetime | None = Field(default=None, description="When session ended") @field_validator("pipeline_phase") @classmethod @@ -194,53 +246,48 @@ def validate_sha(cls, v: Any) -> str | None: return None return str(v) - def get_tool_calls_by_name(self, name: str) -> list[ToolCall]: - """Get all tool calls with the given name.""" - return [tc for tc in self.tool_calls if tc.name == name] - def get_files_by_operation(self, operation: FileOperationType) -> list[FileOperation]: - """Get all file operations of the given type.""" - return [fo for fo in self.files_touched if fo.operation == operation] - - def get_files_written(self) -> list[str]: - """Get paths of all files that were written or created.""" - write_ops = {FileOperationType.WRITE, FileOperationType.CREATE, FileOperationType.EDIT} - return [fo.path for fo in self.files_touched if fo.operation in write_ops] - - def get_files_read(self) -> list[str]: - """Get paths of all files that were read.""" - return [fo.path for fo in self.files_touched if fo.operation == FileOperationType.READ] - - -class CheckpointSummary(BaseModel): - """Summary of a checkpoint for the index.""" +class CheckpointSummaryV2(BaseModel): + """Summary with all queryable fields for the v2 index.""" id: str = Field( ..., pattern=r"^ckpt-[a-f0-9]{8,16}$", description="Unique checkpoint identifier" ) - commit_sha: str = Field(..., pattern=r"^[a-f0-9]{7,40}$", description="Git commit SHA") - session_id: str = Field(..., description="Session ID") - agent_role: str | None = Field(default=None, description="Agent role") + trigger_type: TriggerType = Field(..., description="What created this checkpoint") + session_status: SessionStatus | None = Field(default=None, description="Terminal session state") + + # All queryable fields + session_id: str = Field(..., description="Session/container ID") + commit_sha: str | None = Field( + default=None, pattern=r"^[a-f0-9]{7,40}$", description="Git commit SHA" + ) issue_number: int | None = Field(default=None, ge=1, description="GitHub issue number") pr_number: int | None = Field(default=None, ge=1, description="GitHub PR number") branch: str | None = Field(default=None, description="Git branch") + agent_type: AgentType = Field( + default=AgentType.UNKNOWN, description="Agent role classification" + ) pipeline_phase: str | None = Field(default=None, description="Pipeline phase") + + # Metrics created_at: datetime = Field(..., description="When checkpoint was created") message_count: int = Field(default=0, ge=0, description="Number of messages in transcript") tool_call_count: int = Field(default=0, ge=0, description="Number of tool calls") total_tokens: int = Field(default=0, ge=0, description="Total tokens used") @classmethod - def from_checkpoint(cls, checkpoint: Checkpoint) -> "CheckpointSummary": - """Create a summary from a full checkpoint.""" + def from_checkpoint(cls, checkpoint: "CheckpointV2") -> "CheckpointSummaryV2": + """Create a summary from a full v2 checkpoint.""" return cls( id=checkpoint.id, + trigger_type=checkpoint.trigger_type, + session_status=checkpoint.session_status, + session_id=checkpoint.session_id, commit_sha=checkpoint.commit_sha, - session_id=checkpoint.session.session_id, - agent_role=checkpoint.session.agent_role, issue_number=checkpoint.issue_number, pr_number=checkpoint.pr_number, branch=checkpoint.branch, + agent_type=checkpoint.agent_type, pipeline_phase=checkpoint.pipeline_phase, created_at=checkpoint.created_at, message_count=checkpoint.transcript.message_count if checkpoint.transcript else 0, @@ -249,33 +296,85 @@ def from_checkpoint(cls, checkpoint: Checkpoint) -> "CheckpointSummary": ) -class CheckpointIndex(BaseModel): +class CheckpointIndexV2(BaseModel): """ - Index of checkpoints for a repository. + Multi-dimensional index for fast checkpoint lookups. + + Supports queries like: + - "All checkpoints for issue #530" + - "All checkpoints for PR #42" + - "All checkpoints by session abc123" + - "All checkpoints for commit deadbeef" + - "All coder agent checkpoints in implement phase" + - "All failed sessions" - Stored at the root of the egg/checkpoints/v1 branch to enable - fast lookup of checkpoints by commit SHA, issue number, or branch. + Stored at the root of the egg/checkpoints/v2 branch. """ schemaVersion: str = Field( # noqa: N815 - default="1.0", pattern=r"^[0-9]+\.[0-9]+$", description="Schema version" + default="2.0", pattern=r"^[0-9]+\.[0-9]+$", description="Schema version" ) - checkpoints: list[CheckpointSummary] = Field( + last_updated: datetime = Field(..., description="When the index was last updated") + + # Primary index: all checkpoint summaries + checkpoints: list[CheckpointSummaryV2] = Field( default_factory=list, description="List of checkpoint summaries" ) - last_updated: datetime = Field(..., description="When the index was last updated") - def get_by_commit(self, commit_sha: str) -> CheckpointSummary | None: - """Get checkpoint summary by commit SHA.""" - for cp in self.checkpoints: - if cp.commit_sha == commit_sha or commit_sha.startswith(cp.commit_sha[:7]): - return cp - return None + # Secondary indices for fast lookups (populated on write) + by_session: dict[str, list[str]] = Field( + default_factory=dict, description="session_id -> [checkpoint_ids]" + ) + by_issue: dict[str, list[str]] = Field( + default_factory=dict, description="issue_number (as str) -> [checkpoint_ids]" + ) + by_pr: dict[str, list[str]] = Field( + default_factory=dict, description="pr_number (as str) -> [checkpoint_ids]" + ) + by_commit: dict[str, str] = Field( + default_factory=dict, description="commit_sha -> checkpoint_id (1:1)" + ) + by_agent_type: dict[str, list[str]] = Field( + default_factory=dict, description="agent_type -> [checkpoint_ids]" + ) + by_phase: dict[str, list[str]] = Field( + default_factory=dict, description="pipeline_phase -> [checkpoint_ids]" + ) + by_trigger: dict[str, list[str]] = Field( + default_factory=dict, description="trigger_type -> [checkpoint_ids]" + ) + by_status: dict[str, list[str]] = Field( + default_factory=dict, description="session_status -> [checkpoint_ids]" + ) + + def get_by_session(self, session_id: str) -> list[str]: + """Get checkpoint IDs for a session.""" + return self.by_session.get(session_id, []) + + def get_by_issue(self, issue_number: int) -> list[str]: + """Get checkpoint IDs for an issue.""" + return self.by_issue.get(str(issue_number), []) + + def get_by_pr(self, pr_number: int) -> list[str]: + """Get checkpoint IDs for a PR.""" + return self.by_pr.get(str(pr_number), []) + + def get_by_commit(self, commit_sha: str) -> str | None: + """Get checkpoint ID for a commit (1:1 mapping).""" + return self.by_commit.get(commit_sha) + + def get_by_agent_type(self, agent_type: AgentType) -> list[str]: + """Get checkpoint IDs for an agent type.""" + return self.by_agent_type.get(agent_type.value, []) + + def get_by_phase(self, phase: str) -> list[str]: + """Get checkpoint IDs for a pipeline phase.""" + return self.by_phase.get(phase, []) - def get_by_issue(self, issue_number: int) -> list[CheckpointSummary]: - """Get all checkpoint summaries for an issue.""" - return [cp for cp in self.checkpoints if cp.issue_number == issue_number] + def get_by_trigger(self, trigger_type: TriggerType) -> list[str]: + """Get checkpoint IDs for a trigger type.""" + return self.by_trigger.get(trigger_type.value, []) - def get_by_branch(self, branch: str) -> list[CheckpointSummary]: - """Get all checkpoint summaries for a branch.""" - return [cp for cp in self.checkpoints if cp.branch == branch] + def get_by_status(self, status: SessionStatus) -> list[str]: + """Get checkpoint IDs for a session status.""" + return self.by_status.get(status.value, []) diff --git a/shared/egg_contracts/transcript_extractor.py b/shared/egg_contracts/transcript_extractor.py index 02b5fed9db..2fe0aa3bc9 100644 --- a/shared/egg_contracts/transcript_extractor.py +++ b/shared/egg_contracts/transcript_extractor.py @@ -122,7 +122,7 @@ def extract_session_metadata_from_proxy_buffer( def extract_messages_from_proxy_buffer( entries: list[dict[str, Any]], - max_content_length: int = 10000, + max_content_length: int = 25000, ) -> list[Message]: """ Extract messages from proxy buffer entries. @@ -203,8 +203,8 @@ def extract_messages_from_proxy_buffer( def extract_tool_calls_from_proxy_buffer( entries: list[dict[str, Any]], - max_param_length: int = 1000, - max_result_length: int = 500, + max_param_length: int = 2500, + max_result_length: int = 1500, ) -> tuple[list[ToolCall], list[FileOperation]]: """ Extract tool calls and file operations from proxy buffer entries. @@ -376,9 +376,9 @@ def extract_token_usage_from_proxy_buffer(entries: list[dict[str, Any]]) -> Toke def extract_transcript_from_proxy_buffer( buffer_path: Path, container_id: str | None = None, - max_content_length: int = 10000, - max_param_length: int = 1000, - max_result_length: int = 500, + max_content_length: int = 25000, + max_param_length: int = 2500, + max_result_length: int = 1500, ) -> tuple[SessionMetadata, Transcript, list[ToolCall], list[FileOperation], TokenUsage]: """ Extract full transcript data from a proxy buffer file. diff --git a/shared/egg_contracts/usage.py b/shared/egg_contracts/usage.py index 947321e689..d272f8e461 100644 --- a/shared/egg_contracts/usage.py +++ b/shared/egg_contracts/usage.py @@ -3,7 +3,7 @@ These models define the data structures for tracking token usage across multiple dimensions: session, issue, workflow, and PR. Usage aggregates -are pre-computed JSON files stored in the egg/checkpoints/v1 orphaned +are pre-computed JSON files stored in the egg/checkpoints/v2 orphaned branch alongside checkpoint data. Architecture: diff --git a/shared/egg_contracts/usage_cli.py b/shared/egg_contracts/usage_cli.py index 66991215e2..2da26b3000 100644 --- a/shared/egg_contracts/usage_cli.py +++ b/shared/egg_contracts/usage_cli.py @@ -3,7 +3,7 @@ Usage CLI for querying token usage across sessions, issues, workflows, and PRs. This CLI provides commands to query and display token usage aggregates that -are stored in the egg/checkpoints/v1 branch alongside checkpoint data. +are stored in the egg/checkpoints/v2 branch alongside checkpoint data. Commands: egg-usage summary Show overall usage summary @@ -42,7 +42,7 @@ ) # Checkpoint branch name -CHECKPOINT_BRANCH = "egg/checkpoints/v1" +CHECKPOINT_BRANCH = "egg/checkpoints/v2" def get_repo_path() -> str: @@ -538,9 +538,7 @@ def create_parser() -> argparse.ArgumentParser: backfill_parser.add_argument( "--issue", type=int, help="Associated issue number (helps find checkpoints)" ) - backfill_parser.add_argument( - "--branch", help="Branch name (helps find checkpoints)" - ) + backfill_parser.add_argument("--branch", help="Branch name (helps find checkpoints)") backfill_parser.set_defaults(func=cmd_backfill_pr) return parser diff --git a/shared/egg_contracts/usage_loader.py b/shared/egg_contracts/usage_loader.py index 853f444516..a52d467826 100644 --- a/shared/egg_contracts/usage_loader.py +++ b/shared/egg_contracts/usage_loader.py @@ -26,7 +26,7 @@ from pathlib import Path from typing import TypeVar -from .checkpoints import Checkpoint, TokenUsage +from .checkpoints import CheckpointV2, TokenUsage from .usage import ( CheckpointReference, IssueUsage, @@ -111,9 +111,7 @@ def _atomic_write(path: Path, data: dict) -> None: json_str = json.dumps(data, indent=2, sort_keys=True, default=str) # Write to temp file in same directory (for atomic rename) - fd, temp_path = tempfile.mkstemp( - suffix=".tmp", prefix=".usage_", dir=path.parent - ) + fd, temp_path = tempfile.mkstemp(suffix=".tmp", prefix=".usage_", dir=path.parent) try: with os.fdopen(fd, "w") as f: f.write(json_str) @@ -379,7 +377,7 @@ def _token_usage_to_counts(token_usage: TokenUsage | None) -> TokenCounts: def update_usage_from_checkpoint( base_dir: Path, - checkpoint: Checkpoint, + checkpoint: CheckpointV2, retry_attempts: int = MAX_RETRY_ATTEMPTS, ) -> None: """ @@ -409,7 +407,7 @@ def update_usage_from_checkpoint( token_counts = _token_usage_to_counts(checkpoint.token_usage) checkpoint_ref = CheckpointReference( checkpoint_id=checkpoint.id, - commit_sha=checkpoint.commit_sha, + commit_sha=checkpoint.commit_sha or "no-commit", created_at=checkpoint.created_at, ) @@ -432,7 +430,9 @@ def update_usage_from_checkpoint( break except (UsageLoadError, UsageSaveError) as e: if attempt == retry_attempts - 1: - logger.error(f"Failed to update issue usage after {retry_attempts} attempts: {e}") + logger.error( + f"Failed to update issue usage after {retry_attempts} attempts: {e}" + ) raise time.sleep(RETRY_DELAY_SECONDS * (attempt + 1)) @@ -462,7 +462,7 @@ def update_usage_from_checkpoint( def _update_session_usage( base_dir: Path, - checkpoint: Checkpoint, + checkpoint: CheckpointV2, token_counts: TokenCounts, checkpoint_ref: CheckpointReference, now: datetime, @@ -516,7 +516,7 @@ def _update_session_usage( def _update_issue_usage( base_dir: Path, - checkpoint: Checkpoint, + checkpoint: CheckpointV2, token_counts: TokenCounts, now: datetime, ) -> None: @@ -573,7 +573,7 @@ def _update_issue_usage( def _update_pr_usage( base_dir: Path, - checkpoint: Checkpoint, + checkpoint: CheckpointV2, token_counts: TokenCounts, now: datetime, ) -> None: @@ -626,7 +626,7 @@ def _update_pr_usage( def _update_usage_index( base_dir: Path, - checkpoint: Checkpoint, + checkpoint: CheckpointV2, token_counts: TokenCounts, now: datetime, ) -> None: diff --git a/tests/sandbox/test_entrypoint.py b/tests/sandbox/test_entrypoint.py index 78d59b7588..f990291991 100644 --- a/tests/sandbox/test_entrypoint.py +++ b/tests/sandbox/test_entrypoint.py @@ -750,6 +750,7 @@ def test_config_orchestrator_mode_default(self, monkeypatch): """Default orchestrator mode is None (not in orchestrator mode).""" monkeypatch.delenv("EGG_ORCHESTRATOR_MODE", raising=False) monkeypatch.delenv("EGG_PIPELINE_ID", raising=False) + monkeypatch.delenv("EGG_ORCHESTRATOR_URL", raising=False) config = entrypoint.Config() assert config.is_orchestrator_mode is False diff --git a/tests/shared/egg_contracts/test_checkpoint_loader.py b/tests/shared/egg_contracts/test_checkpoint_loader.py new file mode 100644 index 0000000000..b9370ac618 --- /dev/null +++ b/tests/shared/egg_contracts/test_checkpoint_loader.py @@ -0,0 +1,728 @@ +"""Tests for v2 checkpoint loader functions.""" + +import tempfile +from datetime import UTC, datetime, timedelta +from pathlib import Path + +import pytest +from egg_contracts.checkpoint_loader import ( + CheckpointLoadError, + add_checkpoint_to_index_v2, + generate_checkpoint_id_v2, + get_checkpoint_path, + list_checkpoints_v2, + load_checkpoint_by_commit_v2, + load_checkpoint_by_id_v2, + load_checkpoint_index_v2, + load_checkpoint_v2, + save_checkpoint_index_v2, + save_checkpoint_v2, +) +from egg_contracts.checkpoints import ( + AgentType, + CheckpointIndexV2, + CheckpointV2, + SessionMetadata, + SessionStatus, + TokenUsage, + ToolCall, + Transcript, + TriggerType, +) + + +def _make_v2_checkpoint( + checkpoint_id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, + session_id="container-123", + commit_sha="abc1234567890", + session_status=None, + issue_number=None, + pr_number=None, + agent_type=AgentType.UNKNOWN, + pipeline_phase=None, + branch=None, + now=None, +): + """Helper to create a v2 checkpoint for testing.""" + if now is None: + now = datetime.now(UTC) + session = SessionMetadata(session_id="test-session", started_at=now) + return CheckpointV2( + id=checkpoint_id, + trigger_type=trigger_type, + session_status=session_status, + commit_sha=commit_sha, + session_id=session_id, + issue_number=issue_number, + pr_number=pr_number, + agent_type=agent_type, + pipeline_phase=pipeline_phase, + branch=branch, + session=session, + created_at=now, + session_started_at=now, + ) + + +class TestGenerateCheckpointIdV2: + """Tests for generate_checkpoint_id_v2 function.""" + + def test_determinism(self): + """Test same inputs produce same output.""" + fixed_time = datetime(2025, 6, 15, 12, 0, 0, tzinfo=UTC) + id1 = generate_checkpoint_id_v2("session-1", fixed_time) + id2 = generate_checkpoint_id_v2("session-1", fixed_time) + assert id1 == id2 + + def test_uniqueness_different_session(self): + """Test different session IDs produce different IDs.""" + fixed_time = datetime(2025, 6, 15, 12, 0, 0, tzinfo=UTC) + id1 = generate_checkpoint_id_v2("session-1", fixed_time) + id2 = generate_checkpoint_id_v2("session-2", fixed_time) + assert id1 != id2 + + def test_uniqueness_different_timestamp(self): + """Test different timestamps produce different IDs.""" + time1 = datetime(2025, 6, 15, 12, 0, 0, tzinfo=UTC) + time2 = datetime(2025, 6, 15, 12, 0, 1, tzinfo=UTC) + id1 = generate_checkpoint_id_v2("session-1", time1) + id2 = generate_checkpoint_id_v2("session-1", time2) + assert id1 != id2 + + def test_format(self): + """Test generated ID matches expected format.""" + fixed_time = datetime(2025, 6, 15, 12, 0, 0, tzinfo=UTC) + checkpoint_id = generate_checkpoint_id_v2("session-1", fixed_time) + assert checkpoint_id.startswith("ckpt-") + hex_part = checkpoint_id[5:] + assert len(hex_part) == 16 + assert all(c in "0123456789abcdef" for c in hex_part) + + def test_default_timestamp(self): + """Test that omitting timestamp uses current time.""" + id1 = generate_checkpoint_id_v2("session-1") + id2 = generate_checkpoint_id_v2("session-1") + # With microsecond precision, these should almost always be different + # but both should be valid format + assert id1.startswith("ckpt-") + assert id2.startswith("ckpt-") + + +class TestSaveAndLoadCheckpointV2: + """Tests for save_checkpoint_v2 and load_checkpoint_v2 roundtrip.""" + + def test_roundtrip_commit_checkpoint(self): + """Test saving and loading a commit-triggered v2 checkpoint.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "checkpoint.json" + checkpoint = _make_v2_checkpoint( + issue_number=42, + pr_number=10, + agent_type=AgentType.CODER, + pipeline_phase="implement", + branch="egg/test", + ) + + save_checkpoint_v2(checkpoint, path) + assert path.exists() + + loaded = load_checkpoint_v2(path) + assert loaded.id == checkpoint.id + assert loaded.trigger_type == TriggerType.COMMIT + assert loaded.commit_sha == "abc1234567890" + assert loaded.session_id == "container-123" + assert loaded.issue_number == 42 + assert loaded.pr_number == 10 + assert loaded.agent_type == AgentType.CODER + assert loaded.pipeline_phase == "implement" + assert loaded.branch == "egg/test" + assert loaded.schemaVersion == "2.0" + + def test_roundtrip_session_end_checkpoint(self): + """Test saving and loading a session-end v2 checkpoint.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "checkpoint.json" + checkpoint = _make_v2_checkpoint( + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.EXPIRED, + commit_sha=None, + ) + + save_checkpoint_v2(checkpoint, path) + loaded = load_checkpoint_v2(path) + + assert loaded.trigger_type == TriggerType.SESSION_END + assert loaded.session_status == SessionStatus.EXPIRED + assert loaded.commit_sha is None + + def test_roundtrip_with_transcript_and_usage(self): + """Test roundtrip with transcript and token usage.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "checkpoint.json" + now = datetime.now(UTC) + session = SessionMetadata(session_id="test", started_at=now) + transcript = Transcript( + message_count=10, truncated=True, truncation_reason="size limit" + ) + token_usage = TokenUsage(input_tokens=5000, output_tokens=2000, total_tokens=7000) + + checkpoint = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, + commit_sha="abc1234", + session_id="container-1", + session=session, + transcript=transcript, + token_usage=token_usage, + tool_calls=[ToolCall(name="Bash", parameters={}, timestamp=now)], + created_at=now, + session_started_at=now, + ) + + save_checkpoint_v2(checkpoint, path) + loaded = load_checkpoint_v2(path) + + assert loaded.transcript.message_count == 10 + assert loaded.transcript.truncated is True + assert loaded.token_usage.total_tokens == 7000 + assert len(loaded.tool_calls) == 1 + + def test_load_nonexistent_raises(self): + """Test loading non-existent file raises CheckpointLoadError.""" + with pytest.raises(CheckpointLoadError): + load_checkpoint_v2(Path("/nonexistent/path.json")) + + def test_load_invalid_json_raises(self): + """Test loading invalid JSON raises CheckpointLoadError.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "bad.json" + path.write_text("not valid json {{{") + with pytest.raises(CheckpointLoadError): + load_checkpoint_v2(path) + + def test_save_creates_parent_dirs(self): + """Test save creates parent directories if they don't exist.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "subdir" / "deep" / "checkpoint.json" + checkpoint = _make_v2_checkpoint() + save_checkpoint_v2(checkpoint, path) + assert path.exists() + + +class TestLoadCheckpointIndexV2: + """Tests for load_checkpoint_index_v2 function.""" + + def test_nonexistent_returns_empty(self): + """Test that loading from non-existent file returns empty index.""" + index = load_checkpoint_index_v2(Path("/nonexistent/index.json")) + assert len(index.checkpoints) == 0 + assert index.by_session == {} + assert index.by_issue == {} + + def test_invalid_json_raises(self): + """Test that invalid JSON raises CheckpointLoadError.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "index.json" + path.write_text("not valid json") + with pytest.raises(CheckpointLoadError): + load_checkpoint_index_v2(path) + + def test_roundtrip(self): + """Test save and load of an index.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "index.json" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_session={"session-1": ["ckpt-aa00000001"]}, + ) + save_checkpoint_index_v2(index, path) + + loaded = load_checkpoint_index_v2(path) + assert loaded.by_session == {"session-1": ["ckpt-aa00000001"]} + + +class TestAddCheckpointToIndexV2: + """Tests for add_checkpoint_to_index_v2 function.""" + + def test_add_to_empty_index(self): + """Test adding a checkpoint to a new index.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + checkpoint = _make_v2_checkpoint( + issue_number=42, + pr_number=10, + agent_type=AgentType.CODER, + pipeline_phase="implement", + ) + + index = add_checkpoint_to_index_v2(checkpoint, index_path) + + assert len(index.checkpoints) == 1 + assert index.checkpoints[0].id == checkpoint.id + + def test_secondary_indices_populated(self): + """Test that all secondary indices are populated.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + checkpoint = _make_v2_checkpoint( + checkpoint_id="ckpt-aa00000001", + session_id="container-123", + commit_sha="abc1234567890", + issue_number=42, + pr_number=10, + agent_type=AgentType.CODER, + pipeline_phase="implement", + trigger_type=TriggerType.COMMIT, + ) + + index = add_checkpoint_to_index_v2(checkpoint, index_path) + + # by_session + assert index.get_by_session("container-123") == ["ckpt-aa00000001"] + # by_issue + assert index.get_by_issue(42) == ["ckpt-aa00000001"] + # by_pr + assert index.get_by_pr(10) == ["ckpt-aa00000001"] + # by_commit + assert index.get_by_commit("abc1234567890") == "ckpt-aa00000001" + # by_agent_type + assert index.get_by_agent_type(AgentType.CODER) == ["ckpt-aa00000001"] + # by_phase + assert index.get_by_phase("implement") == ["ckpt-aa00000001"] + # by_trigger + assert index.get_by_trigger(TriggerType.COMMIT) == ["ckpt-aa00000001"] + + def test_session_end_status_index(self): + """Test that session_status is indexed for session-end checkpoints.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + checkpoint = _make_v2_checkpoint( + checkpoint_id="ckpt-bb00000002", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.FAILED, + commit_sha=None, + ) + + index = add_checkpoint_to_index_v2(checkpoint, index_path) + + assert index.get_by_status(SessionStatus.FAILED) == ["ckpt-bb00000002"] + assert index.get_by_trigger(TriggerType.SESSION_END) == ["ckpt-bb00000002"] + + def test_deduplication(self): + """Test that duplicate checkpoint IDs are not added.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + checkpoint = _make_v2_checkpoint() + + index1 = add_checkpoint_to_index_v2(checkpoint, index_path) + index2 = add_checkpoint_to_index_v2(checkpoint, index_path) + + assert len(index1.checkpoints) == 1 + assert len(index2.checkpoints) == 1 + + def test_multiple_checkpoints(self): + """Test adding multiple checkpoints to the index.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + now = datetime.now(UTC) + + cp1 = _make_v2_checkpoint( + checkpoint_id="ckpt-aa00000001", + session_id="session-1", + commit_sha="abc1234567890", + issue_number=42, + now=now, + ) + cp2 = _make_v2_checkpoint( + checkpoint_id="ckpt-bb00000002", + session_id="session-1", + commit_sha="def4567890abc", + issue_number=42, + now=now + timedelta(seconds=1), + ) + cp3 = _make_v2_checkpoint( + checkpoint_id="ckpt-cc00000003", + session_id="session-2", + commit_sha="0123456789abcd", + issue_number=99, + now=now + timedelta(seconds=2), + ) + + add_checkpoint_to_index_v2(cp1, index_path) + add_checkpoint_to_index_v2(cp2, index_path) + index = add_checkpoint_to_index_v2(cp3, index_path) + + assert len(index.checkpoints) == 3 + + # session-1 has 2 checkpoints + assert len(index.get_by_session("session-1")) == 2 + # session-2 has 1 checkpoint + assert len(index.get_by_session("session-2")) == 1 + + # issue 42 has 2 checkpoints + assert len(index.get_by_issue(42)) == 2 + # issue 99 has 1 checkpoint + assert len(index.get_by_issue(99)) == 1 + + def test_optional_fields_not_indexed(self): + """Test that None optional fields don't create index entries.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + checkpoint = _make_v2_checkpoint( + issue_number=None, + pr_number=None, + pipeline_phase=None, + ) + + index = add_checkpoint_to_index_v2(checkpoint, index_path) + + assert index.by_issue == {} + assert index.by_pr == {} + assert index.by_phase == {} + # But by_session and by_trigger should still be populated + assert len(index.by_session) == 1 + assert len(index.by_trigger) == 1 + + def test_multi_dimensional_queries(self): + """Test querying the index from multiple dimensions.""" + with tempfile.TemporaryDirectory() as tmpdir: + index_path = Path(tmpdir) / "index.json" + now = datetime.now(UTC) + + # Coder on issue 42, implement phase + cp1 = _make_v2_checkpoint( + checkpoint_id="ckpt-aa00000001", + session_id="session-1", + commit_sha="aaa1234567890", + issue_number=42, + agent_type=AgentType.CODER, + pipeline_phase="implement", + now=now, + ) + # Tester on issue 42, implement phase + cp2 = _make_v2_checkpoint( + checkpoint_id="ckpt-bb00000002", + session_id="session-2", + commit_sha="bbb1234567890", + issue_number=42, + agent_type=AgentType.TESTER, + pipeline_phase="implement", + now=now + timedelta(seconds=1), + ) + # Coder on issue 99, plan phase + cp3 = _make_v2_checkpoint( + checkpoint_id="ckpt-cc00000003", + session_id="session-3", + commit_sha="ccc1234567890", + issue_number=99, + agent_type=AgentType.CODER, + pipeline_phase="plan", + now=now + timedelta(seconds=2), + ) + # Session-end (failed) + cp4 = _make_v2_checkpoint( + checkpoint_id="ckpt-dd00000004", + session_id="session-4", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.FAILED, + commit_sha=None, + now=now + timedelta(seconds=3), + ) + + add_checkpoint_to_index_v2(cp1, index_path) + add_checkpoint_to_index_v2(cp2, index_path) + add_checkpoint_to_index_v2(cp3, index_path) + index = add_checkpoint_to_index_v2(cp4, index_path) + + # Issue 42: 2 checkpoints + assert len(index.get_by_issue(42)) == 2 + # Coder agents: 2 checkpoints + assert len(index.get_by_agent_type(AgentType.CODER)) == 2 + # Implement phase: 2 checkpoints + assert len(index.get_by_phase("implement")) == 2 + # Commit triggers: 3 checkpoints + assert len(index.get_by_trigger(TriggerType.COMMIT)) == 3 + # Session-end triggers: 1 checkpoint + assert len(index.get_by_trigger(TriggerType.SESSION_END)) == 1 + # Failed status: 1 checkpoint + assert len(index.get_by_status(SessionStatus.FAILED)) == 1 + # Each session has exactly 1 checkpoint + for i in range(1, 5): + assert len(index.get_by_session(f"session-{i}")) == 1 + + +def _save_checkpoint_at_path(checkpoint: CheckpointV2, checkpoints_dir: Path) -> Path: + """Save a checkpoint to the correct path under checkpoints_dir.""" + path = get_checkpoint_path(checkpoints_dir, checkpoint.id) + save_checkpoint_v2(checkpoint, path) + return path + + +class TestLoadCheckpointByIdV2: + """Tests for load_checkpoint_by_id_v2 function.""" + + def test_load_existing(self): + """Test loading a checkpoint by its ID.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir = Path(tmpdir) / "checkpoints" + cp = _make_v2_checkpoint(checkpoint_id="ckpt-aa00000001") + _save_checkpoint_at_path(cp, checkpoints_dir) + + loaded = load_checkpoint_by_id_v2("ckpt-aa00000001", checkpoints_dir) + assert loaded is not None + assert loaded.id == "ckpt-aa00000001" + + def test_load_nonexistent_returns_none(self): + """Test loading a non-existent checkpoint ID returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir = Path(tmpdir) / "checkpoints" + result = load_checkpoint_by_id_v2("ckpt-zz99999999", checkpoints_dir) + assert result is None + + +class TestLoadCheckpointByCommitV2: + """Tests for load_checkpoint_by_commit_v2 function.""" + + def test_load_by_commit(self): + """Test loading a checkpoint by commit SHA using the v2 index.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir = Path(tmpdir) / "checkpoints" + index_path = Path(tmpdir) / "index.json" + + cp = _make_v2_checkpoint( + checkpoint_id="ckpt-aa00000001", + commit_sha="abc1234567890", + ) + _save_checkpoint_at_path(cp, checkpoints_dir) + add_checkpoint_to_index_v2(cp, index_path) + + loaded = load_checkpoint_by_commit_v2("abc1234567890", checkpoints_dir, index_path) + assert loaded is not None + assert loaded.commit_sha == "abc1234567890" + + def test_commit_not_found_returns_none(self): + """Test that a missing commit SHA returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir = Path(tmpdir) / "checkpoints" + index_path = Path(tmpdir) / "index.json" + + result = load_checkpoint_by_commit_v2("deadbeef", checkpoints_dir, index_path) + assert result is None + + def test_index_missing_returns_none(self): + """Test that a missing index file returns None gracefully.""" + with tempfile.TemporaryDirectory() as tmpdir: + result = load_checkpoint_by_commit_v2( + "abc1234", Path(tmpdir) / "checkpoints", Path(tmpdir) / "no-index.json" + ) + assert result is None + + +class TestListCheckpointsV2: + """Tests for list_checkpoints_v2 function.""" + + def _build_index(self, tmpdir: str) -> tuple[Path, Path]: + """Build an index with 4 checkpoints for testing filters.""" + checkpoints_dir = Path(tmpdir) / "checkpoints" + index_path = Path(tmpdir) / "index.json" + now = datetime.now(UTC) + + checkpoints = [ + _make_v2_checkpoint( + checkpoint_id="ckpt-aa00000001", + session_id="session-1", + commit_sha="aaa1234567890", + issue_number=42, + pr_number=10, + agent_type=AgentType.CODER, + pipeline_phase="implement", + branch="egg/feature", + now=now, + ), + _make_v2_checkpoint( + checkpoint_id="ckpt-bb00000002", + session_id="session-2", + commit_sha="bbb1234567890", + issue_number=42, + agent_type=AgentType.TESTER, + pipeline_phase="implement", + branch="egg/feature", + now=now + timedelta(seconds=1), + ), + _make_v2_checkpoint( + checkpoint_id="ckpt-cc00000003", + session_id="session-3", + commit_sha="ccc1234567890", + issue_number=99, + agent_type=AgentType.CODER, + pipeline_phase="plan", + branch="egg/other", + now=now + timedelta(seconds=2), + ), + _make_v2_checkpoint( + checkpoint_id="ckpt-dd00000004", + session_id="session-4", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.FAILED, + commit_sha=None, + issue_number=42, + agent_type=AgentType.CODER, + now=now + timedelta(seconds=3), + ), + ] + + for cp in checkpoints: + _save_checkpoint_at_path(cp, checkpoints_dir) + add_checkpoint_to_index_v2(cp, index_path) + + return checkpoints_dir, index_path + + def test_list_all(self): + """Test listing all checkpoints without filters.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path) + assert len(results) == 4 + + def test_list_sorted_descending(self): + """Test results are sorted by created_at descending.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path) + # Most recent first + assert results[0].id == "ckpt-dd00000004" + assert results[-1].id == "ckpt-aa00000001" + + def test_filter_by_issue(self): + """Test filtering by issue number.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, issue_number=42) + assert len(results) == 3 + assert all(s.issue_number == 42 for s in results) + + def test_filter_by_agent_type(self): + """Test filtering by agent type.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, agent_type="coder") + assert len(results) == 3 + + def test_filter_by_trigger_type(self): + """Test filtering by trigger type.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, trigger_type="session_end") + assert len(results) == 1 + assert results[0].id == "ckpt-dd00000004" + + def test_filter_by_status(self): + """Test filtering by session status.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, session_status="failed") + assert len(results) == 1 + assert results[0].session_status == SessionStatus.FAILED + + def test_filter_by_session_id(self): + """Test filtering by session ID.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, session_id="session-1") + assert len(results) == 1 + assert results[0].id == "ckpt-aa00000001" + + def test_filter_by_branch(self): + """Test filtering by branch name.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, branch="egg/feature") + assert len(results) == 2 + + def test_filter_by_phase(self): + """Test filtering by pipeline phase.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, pipeline_phase="implement") + assert len(results) == 2 + + def test_filter_by_pr(self): + """Test filtering by PR number.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, pr_number=10) + assert len(results) == 1 + assert results[0].id == "ckpt-aa00000001" + + def test_intersect_multiple_filters(self): + """Test that multiple filters are intersected (AND logic).""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + # Issue 42 AND coder agent + results = list_checkpoints_v2( + checkpoints_dir, index_path, issue_number=42, agent_type="coder" + ) + assert len(results) == 2 # ckpt-aa00000001 and ckpt-dd00000004 + + def test_limit(self): + """Test limiting results.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, limit=2) + assert len(results) == 2 + + def test_empty_index(self): + """Test listing from an empty/missing index.""" + with tempfile.TemporaryDirectory() as tmpdir: + results = list_checkpoints_v2(Path(tmpdir) / "checkpoints", Path(tmpdir) / "index.json") + assert results == [] + + def test_no_match_returns_empty(self): + """Test that unmatched filters return empty list.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + results = list_checkpoints_v2(checkpoints_dir, index_path, issue_number=999) + assert results == [] + + def test_invalid_trigger_type_raises(self): + """Test that an invalid trigger_type value raises ValueError.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + with pytest.raises(ValueError): + list_checkpoints_v2(checkpoints_dir, index_path, trigger_type="bogus") + + def test_invalid_session_status_raises(self): + """Test that an invalid session_status value raises ValueError.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + with pytest.raises(ValueError): + list_checkpoints_v2(checkpoints_dir, index_path, session_status="bogus") + + def test_invalid_agent_type_raises(self): + """Test that an invalid agent_type value raises ValueError.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir, index_path = self._build_index(tmpdir) + with pytest.raises(ValueError): + list_checkpoints_v2(checkpoints_dir, index_path, agent_type="bogus") + + +class TestLoadCheckpointByIdV2EdgeCases: + """Tests for edge cases in load_checkpoint_by_id_v2.""" + + def test_malformed_ckpt_id_returns_none(self): + """Test that a malformed ckpt- ID returns None rather than crashing.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir = Path(tmpdir) / "checkpoints" + checkpoints_dir.mkdir() + result = load_checkpoint_by_id_v2("ckpt-garbage-not-hex", checkpoints_dir) + assert result is None + + def test_empty_ckpt_id_returns_none(self): + """Test that 'ckpt-' alone returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + checkpoints_dir = Path(tmpdir) / "checkpoints" + checkpoints_dir.mkdir() + result = load_checkpoint_by_id_v2("ckpt-", checkpoints_dir) + assert result is None diff --git a/tests/shared/egg_contracts/test_checkpoints.py b/tests/shared/egg_contracts/test_checkpoints.py index e513d2baab..af2563d920 100644 --- a/tests/shared/egg_contracts/test_checkpoints.py +++ b/tests/shared/egg_contracts/test_checkpoints.py @@ -2,36 +2,27 @@ from datetime import UTC, datetime, timedelta from pathlib import Path -import tempfile -import json import pytest - +from egg_contracts.checkpoint_loader import ( + generate_checkpoint_id_from_commit, + get_checkpoint_path, +) from egg_contracts.checkpoints import ( - Checkpoint, - CheckpointIndex, - CheckpointSummary, + AgentType, + CheckpointIndexV2, + CheckpointSummaryV2, + CheckpointV2, FileOperation, FileOperationType, Message, MessageRole, SessionMetadata, + SessionStatus, TokenUsage, ToolCall, Transcript, -) -from egg_contracts.checkpoint_loader import ( - CheckpointLoadError, - CheckpointSaveError, - generate_checkpoint_id, - generate_checkpoint_id_from_commit, - get_checkpoint_filename, - get_checkpoint_path, - load_checkpoint, - save_checkpoint, - add_checkpoint_to_index, - load_checkpoint_index, - list_checkpoints, + TriggerType, ) @@ -183,143 +174,310 @@ def test_default_token_usage(self): assert usage.total_tokens == 0 -class TestCheckpoint: - """Tests for Checkpoint model.""" +class TestCheckpointLoaderUtils: + """Tests for shared checkpoint loading utilities.""" + + def test_generate_checkpoint_id_from_commit(self): + """Test deterministic checkpoint ID generation with timestamp.""" + # With same timestamp, same inputs produce same output + fixed_time = datetime(2025, 1, 15, 12, 0, 0, tzinfo=UTC) + id1 = generate_checkpoint_id_from_commit("abc123", "session-1", fixed_time) + id2 = generate_checkpoint_id_from_commit("abc123", "session-1", fixed_time) + id3 = generate_checkpoint_id_from_commit("abc123", "session-2", fixed_time) + + # Same inputs (including timestamp) should produce same output + assert id1 == id2 + + # Different session should produce different output + assert id1 != id3 + + # Different timestamp should produce different output + later_time = datetime(2025, 1, 15, 12, 0, 1, tzinfo=UTC) + id4 = generate_checkpoint_id_from_commit("abc123", "session-1", later_time) + assert id1 != id4 + + # Generated ID should be 16 hex chars (8 bytes) + assert id1.startswith("ckpt-") + hex_part = id1[5:] + assert len(hex_part) == 16 + assert all(c in "0123456789abcdef" for c in hex_part) + + def test_get_checkpoint_path(self): + """Test checkpoint path generation.""" + base = Path("/tmp/checkpoints") + path = get_checkpoint_path(base, "ckpt-ab123456") + + # Should use first 2 chars after prefix as subdirectory + assert path.parent.name == "ab" + assert path.name == "ckpt-ab123456.json" + + +# ============================================================================== +# V2 Model Tests +# ============================================================================== + + +class TestTriggerType: + """Tests for TriggerType enum.""" + + def test_enum_values(self): + """Test TriggerType enum has expected values.""" + assert TriggerType.COMMIT == "commit" + assert TriggerType.SESSION_END == "session_end" + + def test_serialization(self): + """Test TriggerType values can be used as strings.""" + assert TriggerType("commit") == TriggerType.COMMIT + assert TriggerType("session_end") == TriggerType.SESSION_END + + def test_invalid_value(self): + """Test invalid TriggerType raises ValueError.""" + with pytest.raises(ValueError): + TriggerType("invalid") + + +class TestSessionStatus: + """Tests for SessionStatus enum.""" + + def test_enum_values(self): + """Test SessionStatus enum has expected values.""" + assert SessionStatus.COMPLETED == "completed" + assert SessionStatus.EXPIRED == "expired" + assert SessionStatus.FAILED == "failed" + + def test_serialization(self): + """Test SessionStatus values can be used as strings.""" + assert SessionStatus("completed") == SessionStatus.COMPLETED + assert SessionStatus("expired") == SessionStatus.EXPIRED + assert SessionStatus("failed") == SessionStatus.FAILED + + def test_invalid_value(self): + """Test invalid SessionStatus raises ValueError.""" + with pytest.raises(ValueError): + SessionStatus("cancelled") + + +class TestAgentType: + """Tests for AgentType enum.""" + + def test_enum_values(self): + """Test AgentType enum has expected values.""" + assert AgentType.CODER == "coder" + assert AgentType.TESTER == "tester" + assert AgentType.DOCUMENTER == "documenter" + assert AgentType.INTEGRATOR == "integrator" + assert AgentType.REVIEWER == "reviewer" + assert AgentType.UNKNOWN == "unknown" + + def test_serialization(self): + """Test AgentType values can be used as strings.""" + for agent_type in AgentType: + assert AgentType(agent_type.value) == agent_type + + +class TestCheckpointV2: + """Tests for CheckpointV2 model.""" + + def _make_session_metadata(self, now=None): + """Create a SessionMetadata for testing.""" + if now is None: + now = datetime.now(UTC) + return SessionMetadata(session_id="test-session", started_at=now) - def test_minimal_checkpoint(self): - """Test creating a minimal checkpoint.""" + def test_minimal_commit_checkpoint(self): + """Test creating a minimal commit-triggered checkpoint.""" now = datetime.now(UTC) - session = SessionMetadata( - session_id="test-session", - started_at=now, + session = self._make_session_metadata(now) + checkpoint = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, + commit_sha="abc1234567890", + session_id="container-123", + session=session, + created_at=now, + session_started_at=now, ) - checkpoint = Checkpoint( + assert checkpoint.trigger_type == TriggerType.COMMIT + assert checkpoint.commit_sha == "abc1234567890" + assert checkpoint.session_status is None + assert checkpoint.schemaVersion == "2.0" + + def test_session_end_checkpoint_without_commit(self): + """Test creating a session-end checkpoint without commit_sha.""" + now = datetime.now(UTC) + session = self._make_session_metadata(now) + checkpoint = CheckpointV2( id="ckpt-abc123def456", - commit_sha="abc1234", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.COMPLETED, + session_id="container-456", session=session, created_at=now, + session_started_at=now, ) - assert checkpoint.id == "ckpt-abc123def456" - assert checkpoint.commit_sha == "abc1234" - assert checkpoint.transcript is None - assert checkpoint.files_touched == [] + assert checkpoint.trigger_type == TriggerType.SESSION_END + assert checkpoint.commit_sha is None + assert checkpoint.session_status == SessionStatus.COMPLETED def test_full_checkpoint(self): - """Test creating a full checkpoint.""" + """Test creating a full v2 checkpoint with all fields.""" now = datetime.now(UTC) session = SessionMetadata( session_id="test-session", started_at=now, agent_role="coder", + container_id="container-789", ) transcript = Transcript( - messages=[ - Message(role=MessageRole.USER, content="Hello", timestamp=now), - ], + messages=[Message(role=MessageRole.USER, content="Hello", timestamp=now)], message_count=1, ) - tool_calls = [ - ToolCall(name="Bash", parameters={"command": "ls"}, timestamp=now), - ] - files_touched = [ - FileOperation(path="src/main.py", operation=FileOperationType.READ), - ] - token_usage = TokenUsage( - input_tokens=1000, - output_tokens=500, - total_tokens=1500, - ) + token_usage = TokenUsage(input_tokens=1000, output_tokens=500, total_tokens=1500) + tool_calls = [ToolCall(name="Bash", parameters={"command": "ls"}, timestamp=now)] + files_touched = [FileOperation(path="src/main.py", operation=FileOperationType.READ)] - checkpoint = Checkpoint( + checkpoint = CheckpointV2( id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, commit_sha="abc1234567890", + push_sha="def4567890abc", + branch="egg/feature-123", + session_id="container-789", + issue_number=530, + pr_number=42, + agent_type=AgentType.CODER, + pipeline_phase="implement", session=session, transcript=transcript, files_touched=files_touched, tool_calls=tool_calls, token_usage=token_usage, - issue_number=123, - pipeline_phase="implement", - branch="egg/feature-123", created_at=now, + session_started_at=now, + session_ended_at=now + timedelta(hours=1), ) - assert checkpoint.issue_number == 123 + assert checkpoint.issue_number == 530 + assert checkpoint.pr_number == 42 + assert checkpoint.agent_type == AgentType.CODER assert checkpoint.pipeline_phase == "implement" - assert len(checkpoint.tool_calls) == 1 - assert len(checkpoint.files_touched) == 1 + assert checkpoint.session_ended_at is not None - def test_checkpoint_helper_methods(self): - """Test checkpoint helper methods.""" + def test_optional_commit_sha_for_session_end(self): + """Test that commit_sha is truly optional for session-end checkpoints.""" now = datetime.now(UTC) - session = SessionMetadata(session_id="test", started_at=now) - checkpoint = Checkpoint( + session = self._make_session_metadata(now) + checkpoint = CheckpointV2( id="ckpt-abc123def456", - commit_sha="abc1234", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.EXPIRED, + session_id="container-expired", session=session, created_at=now, - tool_calls=[ - ToolCall(name="Bash", parameters={}, timestamp=now), - ToolCall(name="Bash", parameters={}, timestamp=now), - ToolCall(name="Read", parameters={}, timestamp=now), - ], - files_touched=[ - FileOperation(path="a.py", operation=FileOperationType.READ), - FileOperation(path="b.py", operation=FileOperationType.WRITE), - FileOperation(path="c.py", operation=FileOperationType.EDIT), - ], + session_started_at=now, ) + assert checkpoint.commit_sha is None + assert checkpoint.push_sha is None - bash_calls = checkpoint.get_tool_calls_by_name("Bash") - assert len(bash_calls) == 2 - - read_ops = checkpoint.get_files_by_operation(FileOperationType.READ) - assert len(read_ops) == 1 - - written = checkpoint.get_files_written() - assert len(written) == 2 - assert "b.py" in written - assert "c.py" in written + def test_session_id_required(self): + """Test that session_id is required.""" + now = datetime.now(UTC) + session = self._make_session_metadata(now) + with pytest.raises(ValueError): + CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, + commit_sha="abc1234", + # session_id missing + session=session, + created_at=now, + session_started_at=now, + ) - read = checkpoint.get_files_read() - assert len(read) == 1 - assert "a.py" in read + def test_trigger_type_required(self): + """Test that trigger_type is required.""" + now = datetime.now(UTC) + session = self._make_session_metadata(now) + with pytest.raises(ValueError): + CheckpointV2( + id="ckpt-abc123def456", + # trigger_type missing + session_id="container-123", + session=session, + created_at=now, + session_started_at=now, + ) - def test_checkpoint_id_validation(self): - """Test checkpoint ID pattern validation.""" + def test_pipeline_phase_validation(self): + """Test pipeline_phase validates against allowed values.""" now = datetime.now(UTC) - session = SessionMetadata(session_id="test", started_at=now) + session = self._make_session_metadata(now) - # Valid IDs - valid_ids = [ - "ckpt-abc12345", - "ckpt-1234567890ab", - "ckpt-abcdef12", - ] - for cp_id in valid_ids: - checkpoint = Checkpoint( - id=cp_id, + # Valid phases + for phase in ("refine", "plan", "implement", "pr"): + cp = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, + commit_sha="abc1234", + session_id="container-123", + session=session, + created_at=now, + session_started_at=now, + pipeline_phase=phase, + ) + assert cp.pipeline_phase == phase + + # Invalid phase + with pytest.raises(ValueError): + CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, commit_sha="abc1234", + session_id="container-123", session=session, created_at=now, + session_started_at=now, + pipeline_phase="deploy", ) - assert checkpoint.id == cp_id - # Invalid IDs + def test_checkpoint_id_validation(self): + """Test checkpoint ID pattern validation.""" + now = datetime.now(UTC) + session = self._make_session_metadata(now) + with pytest.raises(ValueError): - Checkpoint( + CheckpointV2( id="invalid-id", + trigger_type=TriggerType.COMMIT, commit_sha="abc1234", + session_id="container-123", session=session, created_at=now, + session_started_at=now, ) + def test_empty_commit_sha_becomes_none(self): + """Test that empty string commit_sha is converted to None.""" + now = datetime.now(UTC) + session = self._make_session_metadata(now) + checkpoint = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.COMPLETED, + commit_sha="", + session_id="container-123", + session=session, + created_at=now, + session_started_at=now, + ) + assert checkpoint.commit_sha is None -class TestCheckpointSummary: - """Tests for CheckpointSummary model.""" + +class TestCheckpointSummaryV2: + """Tests for CheckpointSummaryV2 model.""" def test_from_checkpoint(self): - """Test creating summary from checkpoint.""" + """Test creating summary from a full v2 checkpoint.""" now = datetime.now(UTC) session = SessionMetadata( session_id="test-session", @@ -329,9 +487,16 @@ def test_from_checkpoint(self): transcript = Transcript(message_count=5) token_usage = TokenUsage(total_tokens=1500) - checkpoint = Checkpoint( + checkpoint = CheckpointV2( id="ckpt-abc123def456", + trigger_type=TriggerType.COMMIT, commit_sha="abc1234567890", + session_id="container-123", + issue_number=42, + pr_number=10, + branch="egg/test", + agent_type=AgentType.CODER, + pipeline_phase="implement", session=session, transcript=transcript, tool_calls=[ @@ -339,244 +504,158 @@ def test_from_checkpoint(self): ToolCall(name="Read", parameters={}, timestamp=now), ], token_usage=token_usage, - issue_number=42, - branch="egg/test", - pipeline_phase="implement", created_at=now, + session_started_at=now, ) - summary = CheckpointSummary.from_checkpoint(checkpoint) + summary = CheckpointSummaryV2.from_checkpoint(checkpoint) assert summary.id == checkpoint.id - assert summary.commit_sha == checkpoint.commit_sha - assert summary.session_id == "test-session" - assert summary.agent_role == "coder" + assert summary.trigger_type == TriggerType.COMMIT + assert summary.session_status is None + assert summary.session_id == "container-123" + assert summary.commit_sha == "abc1234567890" assert summary.issue_number == 42 + assert summary.pr_number == 10 assert summary.branch == "egg/test" + assert summary.agent_type == AgentType.CODER + assert summary.pipeline_phase == "implement" assert summary.message_count == 5 assert summary.tool_call_count == 2 assert summary.total_tokens == 1500 - -class TestCheckpointIndex: - """Tests for CheckpointIndex model.""" - - def test_empty_index(self): - """Test empty index.""" - now = datetime.now(UTC) - index = CheckpointIndex(last_updated=now) - assert len(index.checkpoints) == 0 - - def test_index_lookups(self): - """Test index lookup methods.""" + def test_from_session_end_checkpoint(self): + """Test creating summary from session-end checkpoint.""" now = datetime.now(UTC) - summaries = [ - CheckpointSummary( - id="ckpt-aa00000001", - commit_sha="abc1234567890", - session_id="session1", - issue_number=100, - branch="egg/feature-100", - created_at=now, - ), - CheckpointSummary( - id="ckpt-bb00000002", - commit_sha="def4567890abc", - session_id="session2", - issue_number=100, - branch="egg/feature-100", - created_at=now, - ), - CheckpointSummary( - id="ckpt-cc00000003", - commit_sha="0123456789abcd", - session_id="session3", - issue_number=200, - branch="egg/feature-200", - created_at=now, - ), - ] - index = CheckpointIndex(last_updated=now, checkpoints=summaries) - - # Test get_by_commit - result = index.get_by_commit("abc1234567890") - assert result is not None - assert result.id == "ckpt-aa00000001" - - # Test partial commit match - result = index.get_by_commit("abc1234") - assert result is not None - assert result.id == "ckpt-aa00000001" - - # Test get_by_issue - issue_100 = index.get_by_issue(100) - assert len(issue_100) == 2 - - issue_200 = index.get_by_issue(200) - assert len(issue_200) == 1 - - # Test get_by_branch - branch_100 = index.get_by_branch("egg/feature-100") - assert len(branch_100) == 2 - - -class TestCheckpointLoader: - """Tests for checkpoint loading/saving utilities.""" - - def test_generate_checkpoint_id(self): - """Test checkpoint ID generation.""" - id1 = generate_checkpoint_id() - id2 = generate_checkpoint_id() - - # Should start with ckpt- - assert id1.startswith("ckpt-") - assert id2.startswith("ckpt-") - - # Should be unique - assert id1 != id2 - - # Should be valid hex after prefix - hex_part = id1[5:] - assert all(c in "0123456789abcdef" for c in hex_part) - - def test_generate_checkpoint_id_from_commit(self): - """Test deterministic checkpoint ID generation with timestamp.""" - # With same timestamp, same inputs produce same output - fixed_time = datetime(2025, 1, 15, 12, 0, 0, tzinfo=UTC) - id1 = generate_checkpoint_id_from_commit("abc123", "session-1", fixed_time) - id2 = generate_checkpoint_id_from_commit("abc123", "session-1", fixed_time) - id3 = generate_checkpoint_id_from_commit("abc123", "session-2", fixed_time) - - # Same inputs (including timestamp) should produce same output - assert id1 == id2 - - # Different session should produce different output - assert id1 != id3 + session = SessionMetadata(session_id="test-session", started_at=now) - # Different timestamp should produce different output - later_time = datetime(2025, 1, 15, 12, 0, 1, tzinfo=UTC) - id4 = generate_checkpoint_id_from_commit("abc123", "session-1", later_time) - assert id1 != id4 - - # Generated ID should be 16 hex chars (8 bytes) - assert id1.startswith("ckpt-") - hex_part = id1[5:] - assert len(hex_part) == 16 - assert all(c in "0123456789abcdef" for c in hex_part) - - def test_get_checkpoint_path(self): - """Test checkpoint path generation.""" - base = Path("/tmp/checkpoints") - path = get_checkpoint_path(base, "ckpt-ab123456") - - # Should use first 2 chars after prefix as subdirectory - assert path.parent.name == "ab" - assert path.name == "ckpt-ab123456.json" - - def test_save_and_load_checkpoint(self): - """Test saving and loading a checkpoint.""" - with tempfile.TemporaryDirectory() as tmpdir: - tmppath = Path(tmpdir) + checkpoint = CheckpointV2( + id="ckpt-abc123def456", + trigger_type=TriggerType.SESSION_END, + session_status=SessionStatus.FAILED, + session_id="container-456", + session=session, + created_at=now, + session_started_at=now, + ) - now = datetime.now(UTC) - session = SessionMetadata(session_id="test", started_at=now) - checkpoint = Checkpoint( - id="ckpt-abc123def456", - commit_sha="abc1234567890", - session=session, - created_at=now, - issue_number=42, - ) + summary = CheckpointSummaryV2.from_checkpoint(checkpoint) - checkpoint_path = get_checkpoint_path(tmppath, checkpoint.id) + assert summary.trigger_type == TriggerType.SESSION_END + assert summary.session_status == SessionStatus.FAILED + assert summary.commit_sha is None + assert summary.message_count == 0 + assert summary.tool_call_count == 0 + assert summary.total_tokens == 0 - # Save - save_checkpoint(checkpoint, checkpoint_path) - assert checkpoint_path.exists() - # Load - loaded = load_checkpoint(checkpoint_path) - assert loaded.id == checkpoint.id - assert loaded.commit_sha == checkpoint.commit_sha - assert loaded.issue_number == 42 +class TestCheckpointIndexV2: + """Tests for CheckpointIndexV2 model.""" - def test_load_nonexistent_checkpoint(self): - """Test loading a non-existent checkpoint.""" - with pytest.raises(CheckpointLoadError): - load_checkpoint(Path("/nonexistent/path.json")) + def _make_summary(self, **kwargs): + """Create a CheckpointSummaryV2 for testing.""" + defaults = { + "id": "ckpt-abc123def456", + "trigger_type": TriggerType.COMMIT, + "session_id": "session-1", + "created_at": datetime.now(UTC), + } + defaults.update(kwargs) + return CheckpointSummaryV2(**defaults) - def test_add_checkpoint_to_index(self): - """Test adding checkpoint to index.""" - with tempfile.TemporaryDirectory() as tmpdir: - tmppath = Path(tmpdir) - index_path = tmppath / "index.json" + def test_empty_index(self): + """Test empty v2 index.""" + now = datetime.now(UTC) + index = CheckpointIndexV2(last_updated=now) + assert len(index.checkpoints) == 0 + assert index.by_session == {} + assert index.by_issue == {} + assert index.by_commit == {} - now = datetime.now(UTC) - session = SessionMetadata(session_id="test", started_at=now) - checkpoint = Checkpoint( - id="ckpt-abc123def456", - commit_sha="abc1234567890", - session=session, - created_at=now, - ) + def test_get_by_session(self): + """Test get_by_session lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_session={"session-1": ["ckpt-aa00000001", "ckpt-aa00000002"]}, + ) + assert index.get_by_session("session-1") == ["ckpt-aa00000001", "ckpt-aa00000002"] + assert index.get_by_session("nonexistent") == [] - # Add to index - index = add_checkpoint_to_index(checkpoint, index_path) + def test_get_by_issue(self): + """Test get_by_issue lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_issue={"42": ["ckpt-aa00000001"]}, + ) + assert index.get_by_issue(42) == ["ckpt-aa00000001"] + assert index.get_by_issue(999) == [] - assert len(index.checkpoints) == 1 - assert index.checkpoints[0].id == checkpoint.id + def test_get_by_pr(self): + """Test get_by_pr lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_pr={"10": ["ckpt-aa00000001", "ckpt-bb00000002"]}, + ) + assert index.get_by_pr(10) == ["ckpt-aa00000001", "ckpt-bb00000002"] + assert index.get_by_pr(999) == [] - # Add another - checkpoint2 = Checkpoint( - id="ckpt-def456789012", - commit_sha="def4567890abc", - session=session, - created_at=now, - ) - index = add_checkpoint_to_index(checkpoint2, index_path) + def test_get_by_commit(self): + """Test get_by_commit lookup (1:1 mapping).""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_commit={"abc1234567890": "ckpt-aa00000001"}, + ) + assert index.get_by_commit("abc1234567890") == "ckpt-aa00000001" + assert index.get_by_commit("nonexistent") is None - assert len(index.checkpoints) == 2 + def test_get_by_agent_type(self): + """Test get_by_agent_type lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_agent_type={"coder": ["ckpt-aa00000001"], "tester": ["ckpt-bb00000002"]}, + ) + assert index.get_by_agent_type(AgentType.CODER) == ["ckpt-aa00000001"] + assert index.get_by_agent_type(AgentType.TESTER) == ["ckpt-bb00000002"] + assert index.get_by_agent_type(AgentType.UNKNOWN) == [] - def test_list_checkpoints_empty(self): - """Test listing checkpoints from empty directory.""" - with tempfile.TemporaryDirectory() as tmpdir: - tmppath = Path(tmpdir) - checkpoints = list_checkpoints(tmppath) - assert checkpoints == [] + def test_get_by_phase(self): + """Test get_by_phase lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_phase={"implement": ["ckpt-aa00000001"]}, + ) + assert index.get_by_phase("implement") == ["ckpt-aa00000001"] + assert index.get_by_phase("plan") == [] - def test_list_checkpoints_with_filter(self): - """Test listing checkpoints with filters.""" - with tempfile.TemporaryDirectory() as tmpdir: - tmppath = Path(tmpdir) + def test_get_by_trigger(self): + """Test get_by_trigger lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_trigger={ + "commit": ["ckpt-aa00000001"], + "session_end": ["ckpt-bb00000002"], + }, + ) + assert index.get_by_trigger(TriggerType.COMMIT) == ["ckpt-aa00000001"] + assert index.get_by_trigger(TriggerType.SESSION_END) == ["ckpt-bb00000002"] - now = datetime.now(UTC) - session = SessionMetadata(session_id="test", started_at=now) - - # Create some checkpoints - for i in range(3): - cp = Checkpoint( - id=f"ckpt-aa{i:010d}", - commit_sha=f"abcdef{i:034d}", - session=session, - created_at=now - timedelta(hours=i), - issue_number=100 if i < 2 else 200, - branch="egg/test-100" if i < 2 else "egg/test-200", - ) - path = get_checkpoint_path(tmppath, cp.id) - save_checkpoint(cp, path) - - # List all - all_cps = list_checkpoints(tmppath) - assert len(all_cps) == 3 - - # Filter by issue - issue_100 = list_checkpoints(tmppath, issue_number=100) - assert len(issue_100) == 2 - - # Filter by branch - branch_200 = list_checkpoints(tmppath, branch="egg/test-200") - assert len(branch_200) == 1 - - # With limit - limited = list_checkpoints(tmppath, limit=2) - assert len(limited) == 2 + def test_get_by_status(self): + """Test get_by_status lookup.""" + now = datetime.now(UTC) + index = CheckpointIndexV2( + last_updated=now, + by_status={ + "completed": ["ckpt-aa00000001"], + "failed": ["ckpt-bb00000002"], + }, + ) + assert index.get_by_status(SessionStatus.COMPLETED) == ["ckpt-aa00000001"] + assert index.get_by_status(SessionStatus.FAILED) == ["ckpt-bb00000002"] + assert index.get_by_status(SessionStatus.EXPIRED) == [] diff --git a/tests/shared/egg_contracts/test_redactor.py b/tests/shared/egg_contracts/test_redactor.py index c62360aea8..3768239869 100644 --- a/tests/shared/egg_contracts/test_redactor.py +++ b/tests/shared/egg_contracts/test_redactor.py @@ -1,14 +1,12 @@ """Tests for the sensitive data redactor.""" -import pytest - from egg_contracts.redactor import ( + REDACTED_PLACEHOLDER, Redactor, RedactorConfig, - redact_text, - redact_dict, redact_command, - REDACTED_PLACEHOLDER, + redact_dict, + redact_text, ) diff --git a/tests/shared/egg_contracts/test_transcript_extractor.py b/tests/shared/egg_contracts/test_transcript_extractor.py index f2b4d84aad..d7202919d2 100644 --- a/tests/shared/egg_contracts/test_transcript_extractor.py +++ b/tests/shared/egg_contracts/test_transcript_extractor.py @@ -74,9 +74,7 @@ def test_extract_user_message(self): { "timestamp": "2026-02-11T10:00:00.000000Z", "type": "api_turn", - "request": { - "messages": [{"role": "user", "content": "Hello, Claude!"}] - }, + "request": {"messages": [{"role": "user", "content": "Hello, Claude!"}]}, } ] @@ -93,9 +91,7 @@ def test_extract_assistant_text_response(self): "timestamp": "2026-02-11T10:00:00.000000Z", "type": "api_turn", "request": {"messages": []}, - "response": { - "content": [{"type": "text", "text": "Hello! How can I help?"}] - }, + "response": {"content": [{"type": "text", "text": "Hello! How can I help?"}]}, } ] diff --git a/tests/shared/egg_contracts/test_usage_loader.py b/tests/shared/egg_contracts/test_usage_loader.py index e68aebceb8..f7028db87f 100644 --- a/tests/shared/egg_contracts/test_usage_loader.py +++ b/tests/shared/egg_contracts/test_usage_loader.py @@ -6,9 +6,10 @@ import pytest from egg_contracts.checkpoints import ( - Checkpoint, + CheckpointV2, SessionMetadata, TokenUsage, + TriggerType, ) from egg_contracts.usage import ( IssueUsage, @@ -232,12 +233,14 @@ def _create_test_checkpoint( pr_number: int | None = None, input_tokens: int = 1000, output_tokens: int = 500, - ) -> Checkpoint: + ) -> CheckpointV2: """Create a test checkpoint.""" now = datetime.now(UTC) - return Checkpoint( + return CheckpointV2( id="ckpt-abc123456789", + trigger_type=TriggerType.COMMIT, commit_sha="abc123456789", + session_id=session_id, session=SessionMetadata( session_id=session_id, container_id="container-123", @@ -255,6 +258,7 @@ def _create_test_checkpoint( branch="egg/issue-519", pipeline_phase="implement", created_at=now, + session_started_at=now, ) def test_update_creates_new_session_usage(self): diff --git a/tests/shared/test_egg_orchestrator.py b/tests/shared/test_egg_orchestrator.py index ac0384f1b8..3ffc8aeb5a 100644 --- a/tests/shared/test_egg_orchestrator.py +++ b/tests/shared/test_egg_orchestrator.py @@ -165,11 +165,13 @@ class TestSignalResponse: def test_from_dict_success(self): """Test parsing successful response.""" - response = SignalResponse.from_dict({ - "success": True, - "message": "Signal received", - "data": {"id": 123}, - }) + response = SignalResponse.from_dict( + { + "success": True, + "message": "Signal received", + "data": {"id": 123}, + } + ) assert response.success is True assert response.message == "Signal received" @@ -191,6 +193,7 @@ def test_is_orchestrator_mode_false(self, monkeypatch): """Test detection when not in orchestrator mode.""" monkeypatch.delenv("EGG_ORCHESTRATOR_MODE", raising=False) monkeypatch.delenv("EGG_PIPELINE_ID", raising=False) + monkeypatch.delenv("EGG_ORCHESTRATOR_URL", raising=False) assert is_orchestrator_mode() is False def test_is_orchestrator_mode_true(self, monkeypatch): @@ -203,6 +206,7 @@ def test_get_orchestrator_url_not_set(self, monkeypatch): """Test URL detection when not in orchestrator mode.""" monkeypatch.delenv("EGG_ORCHESTRATOR_URL", raising=False) monkeypatch.delenv("EGG_ORCHESTRATOR_MODE", raising=False) + monkeypatch.delenv("EGG_PIPELINE_ID", raising=False) assert get_orchestrator_url() is None def test_get_orchestrator_url_explicit(self, monkeypatch): @@ -233,6 +237,7 @@ def test_singleton_returns_same_instance(self, monkeypatch): """Test that singleton returns the same instance.""" # Reset the global client first import egg_orchestrator.client as client_module + client_module._orchestrator_client = None monkeypatch.setenv("EGG_ORCHESTRATOR_URL", "http://test:8080") @@ -270,11 +275,13 @@ def _create_mock_response(self, data: dict) -> MagicMock: def test_signal_complete_success(self, client, mock_urlopen): """Test signal_complete sends correct HTTP request.""" - mock_urlopen.return_value = self._create_mock_response({ - "success": True, - "message": "Signal received", - "data": {"signal_id": 123}, - }) + mock_urlopen.return_value = self._create_mock_response( + { + "success": True, + "message": "Signal received", + "data": {"signal_id": 123}, + } + ) response = client.signal_complete( pipeline_id="issue-456", @@ -300,10 +307,12 @@ def test_signal_complete_success(self, client, mock_urlopen): def test_signal_error_success(self, client, mock_urlopen): """Test signal_error sends correct HTTP request.""" - mock_urlopen.return_value = self._create_mock_response({ - "success": True, - "message": "Error signal received", - }) + mock_urlopen.return_value = self._create_mock_response( + { + "success": True, + "message": "Error signal received", + } + ) response = client.signal_error( pipeline_id="issue-789", @@ -326,10 +335,12 @@ def test_signal_error_success(self, client, mock_urlopen): def test_signal_progress_success(self, client, mock_urlopen): """Test signal_progress sends correct HTTP request.""" - mock_urlopen.return_value = self._create_mock_response({ - "success": True, - "message": "Progress updated", - }) + mock_urlopen.return_value = self._create_mock_response( + { + "success": True, + "message": "Progress updated", + } + ) response = client.signal_progress( pipeline_id="issue-100", @@ -352,10 +363,12 @@ def test_signal_progress_success(self, client, mock_urlopen): def test_signal_heartbeat_success(self, client, mock_urlopen): """Test signal_heartbeat sends correct HTTP request.""" - mock_urlopen.return_value = self._create_mock_response({ - "success": True, - "message": "Heartbeat received", - }) + mock_urlopen.return_value = self._create_mock_response( + { + "success": True, + "message": "Heartbeat received", + } + ) response = client.signal_heartbeat( pipeline_id="issue-200", @@ -379,10 +392,12 @@ def test_signal_http_error_with_json_response(self, client, mock_urlopen): from egg_orchestrator.client import OrchestratorError error_response = MagicMock() - error_response.read.return_value = json.dumps({ - "message": "Pipeline not found", - "details": {"pipeline_id": "invalid"}, - }).encode() + error_response.read.return_value = json.dumps( + { + "message": "Pipeline not found", + "details": {"pipeline_id": "invalid"}, + } + ).encode() mock_urlopen.side_effect = HTTPError( url="http://test",