from __future__ import annotations import asyncio from collections.abc import Sequence from dataclasses import dataclass from datetime import datetime from unittest.mock import MagicMock import pytest from vibe.core.nuage.workflow import ( WorkflowExecutionListResponse, WorkflowExecutionStatus, WorkflowExecutionWithoutResultResponse, ) from vibe.core.session.resume_sessions import ( RemoteResumeResult, RemoteResumeSessions, ResumeSessionInfo, can_delete_resume_session_source, list_remote_resume_sessions, session_latest_messages, short_session_id, ) from vibe.core.session.session_id import shorten_session_id @dataclass(frozen=True) class RemoteResumeRequest: workflow_identifier: str | None page_size: int status: Sequence[WorkflowExecutionStatus] | None class FakeRemoteResumeClient: def __init__( self, response: WorkflowExecutionListResponse | None = None, *, delay: float = 0.0, ) -> None: self.response = response or WorkflowExecutionListResponse(executions=[]) self.delay = delay self.requests: list[RemoteResumeRequest] = [] self.closed = False async def get_workflow_runs( self, workflow_identifier: str | None = None, page_size: int = 50, next_page_token: str | None = None, status: Sequence[WorkflowExecutionStatus] | None = None, user_id: str = "current", ) -> WorkflowExecutionListResponse: self.requests.append( RemoteResumeRequest( workflow_identifier=workflow_identifier, page_size=page_size, status=status, ) ) if self.delay: await asyncio.sleep(self.delay) return self.response async def aclose(self) -> None: self.closed = True def enabled_vibe_code_config() -> MagicMock: config = MagicMock() config.vibe_code_enabled = True config.vibe_code_api_key = "test-key" config.vibe_code_base_url = "https://test.example.com" config.api_timeout = 30 config.vibe_code_workflow_id = "workflow-1" return config class TestShortenSessionId: def test_shortens_to_first_8_chars(self) -> None: sid = "abcdef1234567890" assert shorten_session_id(sid) == "abcdef12" def test_from_end_shortens_to_last_8_chars(self) -> None: sid = "abcdef1234567890" assert shorten_session_id(sid, from_end=True) == "34567890" def test_returns_full_id_when_shorter_than_limit(self) -> None: sid = "abc" assert shorten_session_id(sid) == "abc" assert shorten_session_id(sid, from_end=True) == "abc" class TestShortSessionId: def test_local_delegates_to_shorten(self) -> None: sid = "abcdef1234567890" assert short_session_id(sid) == shorten_session_id(sid) def test_local_is_default(self) -> None: sid = "abcdef1234567890" assert short_session_id(sid) == short_session_id(sid, source="local") def test_remote_delegates_to_shorten_from_end(self) -> None: sid = "abcdef1234567890" assert short_session_id(sid, source="remote") == shorten_session_id( sid, from_end=True ) def test_empty_string(self) -> None: assert short_session_id("") == "" class TestCanDeleteResumeSession: def test_local_source_can_delete(self) -> None: assert can_delete_resume_session_source("local") is True def test_remote_source_cannot_delete(self) -> None: assert can_delete_resume_session_source("remote") is False def test_session_info_can_delete_matches_source(self) -> None: session = ResumeSessionInfo( session_id="session-a", source="local", cwd="/test", title=None, end_time=None, ) assert session.can_delete is True class TestListRemoteResumeSessions: @pytest.mark.asyncio async def test_passes_active_statuses_to_api(self) -> None: running = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-running", status=WorkflowExecutionStatus.RUNNING, start_time=datetime(2026, 1, 1), end_time=None, ) continued = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-continued", status=WorkflowExecutionStatus.CONTINUED_AS_NEW, start_time=datetime(2026, 1, 1), end_time=None, ) mock_response = WorkflowExecutionListResponse(executions=[running, continued]) client = FakeRemoteResumeClient(mock_response) result = await list_remote_resume_sessions(client, "workflow-1") assert len(result) == 2 session_ids = {s.session_id for s in result} assert "exec-running" in session_ids assert "exec-continued" in session_ids assert all(s.source == "remote" for s in result) assert client.requests == [ RemoteResumeRequest( workflow_identifier="workflow-1", page_size=50, status=[ WorkflowExecutionStatus.RUNNING, WorkflowExecutionStatus.CONTINUED_AS_NEW, ], ) ] @pytest.mark.asyncio async def test_deduplicates_execution_ids_keeps_latest(self) -> None: older = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-1", status=WorkflowExecutionStatus.RUNNING, start_time=datetime(2026, 1, 1), end_time=None, ) newer = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-1", status=WorkflowExecutionStatus.RUNNING, start_time=datetime(2026, 1, 5), end_time=None, ) other = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-2", status=WorkflowExecutionStatus.RUNNING, start_time=datetime(2026, 1, 3), end_time=None, ) mock_response = WorkflowExecutionListResponse(executions=[older, newer, other]) client = FakeRemoteResumeClient(mock_response) result = await list_remote_resume_sessions(client, "workflow-1") assert len(result) == 2 by_id = {s.session_id: s for s in result} assert by_id["exec-1"].end_time == datetime(2026, 1, 5).isoformat() assert "exec-2" in by_id @pytest.mark.asyncio async def test_dedup_keeps_latest_start_time_when_previous_has_end_time( self, ) -> None: previous = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-1", status=WorkflowExecutionStatus.FAILED, start_time=datetime(2026, 1, 1), end_time=datetime(2026, 1, 10), ) newer = WorkflowExecutionWithoutResultResponse( workflow_name="vibe", execution_id="exec-1", status=WorkflowExecutionStatus.RUNNING, start_time=datetime(2026, 1, 5), end_time=None, ) mock_response = WorkflowExecutionListResponse(executions=[previous, newer]) client = FakeRemoteResumeClient(mock_response) result = await list_remote_resume_sessions(client, "workflow-1") assert len(result) == 1 assert result[0].session_id == "exec-1" assert result[0].status == WorkflowExecutionStatus.RUNNING class TestSessionLatestMessages: def test_remote_session_formats_title_and_status(self) -> None: session = ResumeSessionInfo( session_id="exec-1", source="remote", cwd="", title="My run", end_time=None, status="RUNNING", ) messages = session_latest_messages([session], MagicMock()) assert messages[session.option_id] == "My run (running)" class TestRemoteResumeSessions: @pytest.mark.asyncio async def test_fetch_skips_when_vibe_code_disabled(self) -> None: config = MagicMock() config.vibe_code_enabled = False config.vibe_code_api_key = "key" created_clients: list[FakeRemoteResumeClient] = [] def client_factory(_config: object) -> FakeRemoteResumeClient: client = FakeRemoteResumeClient() created_clients.append(client) return client remote = RemoteResumeSessions(lambda: config, client_factory) result = await remote.fetch(10.0) assert result == RemoteResumeResult([], None) assert created_clients == [] @pytest.mark.asyncio async def test_start_cancels_previous_fetch(self) -> None: config = enabled_vibe_code_config() client = FakeRemoteResumeClient(delay=10.0) remote = RemoteResumeSessions(lambda: config, lambda _config: client) first = remote.start(100.0) await asyncio.sleep(0) second = remote.start(100.0) await asyncio.sleep(0) await remote.aclose() assert first.cancelled() assert first is not second @pytest.mark.asyncio async def test_fetch_returns_error_tuple_on_timeout(self) -> None: config = enabled_vibe_code_config() client = FakeRemoteResumeClient(delay=10.0) remote = RemoteResumeSessions(lambda: config, lambda _config: client) sessions, error = await remote.fetch(0.01) await remote.aclose() assert sessions == [] assert error is not None and "Timed out" in error @pytest.mark.asyncio async def test_aclose_cancels_inflight_fetch_before_closing_client(self) -> None: config = enabled_vibe_code_config() client = FakeRemoteResumeClient(delay=10.0) remote = RemoteResumeSessions(lambda: config, lambda _config: client) task = remote.start(100.0) await asyncio.sleep(0) await remote.aclose() assert task.cancelled() assert client.closed is True