v2.7.3 (#564)
Co-authored-by: Bastien <bastien.baret@gmail.com> Co-authored-by: Laure Hugo <201583486+laure0303@users.noreply.github.com> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Paul Cacheux <paul.cacheux@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin
GitHub
parent
9c1c32e058
commit
90763daf81
61 changed files with 6046 additions and 694 deletions
0
tests/core/nuage/__init__.py
Normal file
0
tests/core/nuage/__init__.py
Normal file
283
tests/core/nuage/test_remote_events_source.py
Normal file
283
tests/core/nuage/test_remote_events_source.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.conftest import build_test_vibe_config
|
||||
from vibe.core.agent_loop import AgentLoopStateError
|
||||
from vibe.core.nuage.exceptions import ErrorCode, WorkflowsException
|
||||
from vibe.core.nuage.remote_events_source import RemoteEventsSource
|
||||
from vibe.core.nuage.streaming import StreamEvent
|
||||
|
||||
_SESSION_ID = "test-session"
|
||||
|
||||
|
||||
def _make_source(**kwargs) -> RemoteEventsSource:
|
||||
config = build_test_vibe_config(enabled_tools=kwargs.pop("enabled_tools", []))
|
||||
return RemoteEventsSource(session_id=_SESSION_ID, config=config, **kwargs)
|
||||
|
||||
|
||||
def _make_retryable_exc(msg: str) -> WorkflowsException:
|
||||
return WorkflowsException(message=msg, code=ErrorCode.GET_EVENTS_STREAM_ERROR)
|
||||
|
||||
|
||||
def _make_stream_event(
|
||||
broker_sequence: int | None = None, data: dict | None = None
|
||||
) -> StreamEvent:
|
||||
return StreamEvent(data=data or {}, broker_sequence=broker_sequence)
|
||||
|
||||
|
||||
class TestIsRetryableStreamDisconnect:
|
||||
def test_peer_closed_connection(self) -> None:
|
||||
source = _make_source()
|
||||
exc = _make_retryable_exc("Peer closed connection without response")
|
||||
assert source._is_retryable_stream_disconnect(exc) is True
|
||||
|
||||
def test_incomplete_chunked_read(self) -> None:
|
||||
source = _make_source()
|
||||
exc = _make_retryable_exc("Incomplete chunked read during streaming")
|
||||
assert source._is_retryable_stream_disconnect(exc) is True
|
||||
|
||||
def test_non_retryable_message(self) -> None:
|
||||
source = _make_source()
|
||||
exc = WorkflowsException(
|
||||
message="some other error", code=ErrorCode.GET_EVENTS_STREAM_ERROR
|
||||
)
|
||||
assert source._is_retryable_stream_disconnect(exc) is False
|
||||
|
||||
def test_wrong_error_code(self) -> None:
|
||||
source = _make_source()
|
||||
exc = WorkflowsException(
|
||||
message="peer closed connection",
|
||||
code=ErrorCode.POST_EXECUTIONS_SIGNALS_ERROR,
|
||||
)
|
||||
assert source._is_retryable_stream_disconnect(exc) is False
|
||||
|
||||
|
||||
async def _async_gen_from_list(items):
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
|
||||
async def _async_gen_raise(items, exc):
|
||||
for item in items:
|
||||
yield item
|
||||
raise exc
|
||||
|
||||
|
||||
class _FakeStream:
|
||||
def __init__(self, payloads=None, exc=None):
|
||||
self._payloads = payloads or []
|
||||
self._exc = exc
|
||||
self._closed = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self._iterate().__aiter__()
|
||||
|
||||
async def _iterate(self):
|
||||
for p in self._payloads:
|
||||
yield p
|
||||
if self._exc is not None:
|
||||
raise self._exc
|
||||
|
||||
async def aclose(self):
|
||||
self._closed = True
|
||||
|
||||
|
||||
class TestStreamRemoteEventsRetry:
|
||||
@pytest.mark.asyncio
|
||||
async def test_retries_on_retryable_disconnect(self) -> None:
|
||||
source = _make_source()
|
||||
exc = _make_retryable_exc("peer closed connection")
|
||||
|
||||
call_count = 0
|
||||
|
||||
def make_stream(_params):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return _FakeStream(exc=exc)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = make_stream
|
||||
source._client = mock_client
|
||||
|
||||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||||
events = [e async for e in source._stream_remote_events()]
|
||||
|
||||
assert events == []
|
||||
assert call_count == 4 # 1 initial + 3 retries
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stops_after_max_retry_count(self) -> None:
|
||||
source = _make_source()
|
||||
exc = _make_retryable_exc("incomplete chunked read")
|
||||
|
||||
call_count = 0
|
||||
|
||||
def make_stream(_params):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return _FakeStream(exc=exc)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = make_stream
|
||||
source._client = mock_client
|
||||
|
||||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||||
events = [e async for e in source._stream_remote_events()]
|
||||
|
||||
assert events == []
|
||||
assert call_count == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resets_retry_count_on_successful_event(self) -> None:
|
||||
source = _make_source()
|
||||
exc = _make_retryable_exc("peer closed connection")
|
||||
|
||||
successful_event = _make_stream_event(broker_sequence=0, data={})
|
||||
call_count = 0
|
||||
|
||||
def make_stream(_params):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count <= 2:
|
||||
return _FakeStream(payloads=[successful_event], exc=exc)
|
||||
return _FakeStream(exc=exc)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = make_stream
|
||||
source._client = mock_client
|
||||
|
||||
with (
|
||||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||||
patch.object(source, "_normalize_stream_event", return_value=None),
|
||||
):
|
||||
events = [e async for e in source._stream_remote_events()]
|
||||
|
||||
assert events == []
|
||||
# call 1: success + exc -> retry_count = 1
|
||||
# call 2: success (reset) + exc -> retry_count = 1
|
||||
# call 3: exc -> retry_count = 2
|
||||
# call 4: exc -> retry_count = 3
|
||||
# call 5: exc -> retry_count = 4 > 3 -> break
|
||||
assert call_count == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_retryable_raises_agent_loop_state_error(self) -> None:
|
||||
source = _make_source()
|
||||
exc = WorkflowsException(
|
||||
message="something bad", code=ErrorCode.TEMPORAL_CONNECTION_ERROR
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = lambda _: _FakeStream(exc=exc)
|
||||
source._client = mock_client
|
||||
|
||||
with pytest.raises(AgentLoopStateError):
|
||||
async for _ in source._stream_remote_events():
|
||||
pass
|
||||
|
||||
|
||||
class TestStreamRemoteEventsIdleBoundary:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stops_on_idle_boundary(self) -> None:
|
||||
source = _make_source()
|
||||
event_data = _make_stream_event(broker_sequence=0, data={})
|
||||
sentinel_event = MagicMock()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = lambda _: _FakeStream(payloads=[event_data])
|
||||
source._client = mock_client
|
||||
|
||||
workflow_event = MagicMock()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
source, "_normalize_stream_event", return_value=workflow_event
|
||||
),
|
||||
patch.object(
|
||||
source, "_consume_workflow_event", return_value=[sentinel_event]
|
||||
),
|
||||
patch.object(source, "_is_idle_boundary", return_value=True) as mock_idle,
|
||||
):
|
||||
events = [
|
||||
e
|
||||
async for e in source._stream_remote_events(stop_on_idle_boundary=True)
|
||||
]
|
||||
|
||||
assert events == [sentinel_event]
|
||||
mock_idle.assert_called_once_with(workflow_event)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_continues_past_idle_boundary_when_disabled(self) -> None:
|
||||
source = _make_source()
|
||||
event1 = _make_stream_event(broker_sequence=0, data={})
|
||||
event2 = _make_stream_event(broker_sequence=1, data={})
|
||||
sentinel1 = MagicMock()
|
||||
sentinel2 = MagicMock()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = lambda _: _FakeStream(payloads=[event1, event2])
|
||||
source._client = mock_client
|
||||
|
||||
workflow_event = MagicMock()
|
||||
call_count = 0
|
||||
|
||||
def consume_side_effect(_evt):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return [sentinel1] if call_count == 1 else [sentinel2]
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
source, "_normalize_stream_event", return_value=workflow_event
|
||||
),
|
||||
patch.object(
|
||||
source, "_consume_workflow_event", side_effect=consume_side_effect
|
||||
),
|
||||
patch.object(source, "_is_idle_boundary", return_value=True),
|
||||
):
|
||||
events = [
|
||||
e
|
||||
async for e in source._stream_remote_events(stop_on_idle_boundary=False)
|
||||
]
|
||||
|
||||
assert events == [sentinel1, sentinel2]
|
||||
|
||||
|
||||
class TestBrokerSequenceTracking:
|
||||
@pytest.mark.asyncio
|
||||
async def test_next_start_seq_updated(self) -> None:
|
||||
source = _make_source()
|
||||
assert source._next_start_seq == 0
|
||||
|
||||
event1 = _make_stream_event(broker_sequence=5, data={})
|
||||
event2 = _make_stream_event(broker_sequence=10, data={})
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = lambda _: _FakeStream(payloads=[event1, event2])
|
||||
source._client = mock_client
|
||||
|
||||
with patch.object(source, "_normalize_stream_event", return_value=None):
|
||||
events = [e async for e in source._stream_remote_events()]
|
||||
|
||||
assert events == []
|
||||
assert source._next_start_seq == 11
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_broker_sequence_not_updated(self) -> None:
|
||||
source = _make_source()
|
||||
source._next_start_seq = 5
|
||||
|
||||
event = _make_stream_event(broker_sequence=None, data={})
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.stream_events = lambda _: _FakeStream(payloads=[event])
|
||||
source._client = mock_client
|
||||
|
||||
with patch.object(source, "_normalize_stream_event", return_value=None):
|
||||
events = [e async for e in source._stream_remote_events()]
|
||||
|
||||
assert events == []
|
||||
assert source._next_start_seq == 5
|
||||
206
tests/core/nuage/test_workflows_client.py
Normal file
206
tests/core/nuage/test_workflows_client.py
Normal file
|
|
@ -0,0 +1,206 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.nuage.client import WorkflowsClient
|
||||
from vibe.core.nuage.exceptions import ErrorCode, WorkflowsException
|
||||
from vibe.core.nuage.streaming import StreamEvent, StreamEventsQueryParams
|
||||
|
||||
|
||||
def _make_client() -> WorkflowsClient:
|
||||
return WorkflowsClient(base_url="http://localhost:8080", api_key="test-key")
|
||||
|
||||
|
||||
def _valid_event_payload() -> dict:
|
||||
return {
|
||||
"stream": "test-stream",
|
||||
"timestamp_unix_nano": 1000000,
|
||||
"data": {"key": "value"},
|
||||
}
|
||||
|
||||
|
||||
class TestParseSSEData:
|
||||
def test_valid_json_returns_stream_event(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
result = client._parse_sse_data(json.dumps(payload), event_type=None)
|
||||
assert isinstance(result, StreamEvent)
|
||||
assert result.stream == "test-stream"
|
||||
assert result.data == {"key": "value"}
|
||||
|
||||
def test_error_event_type_raises(self) -> None:
|
||||
client = _make_client()
|
||||
payload = {"some": "data"}
|
||||
with pytest.raises(WorkflowsException) as exc_info:
|
||||
client._parse_sse_data(json.dumps(payload), event_type="error")
|
||||
assert exc_info.value.code == ErrorCode.GET_EVENTS_STREAM_ERROR
|
||||
assert "Stream error from server" in exc_info.value.message
|
||||
|
||||
def test_error_key_in_json_raises(self) -> None:
|
||||
client = _make_client()
|
||||
payload = {"error": "something went wrong"}
|
||||
with pytest.raises(WorkflowsException) as exc_info:
|
||||
client._parse_sse_data(json.dumps(payload), event_type=None)
|
||||
assert exc_info.value.code == ErrorCode.GET_EVENTS_STREAM_ERROR
|
||||
assert "something went wrong" in exc_info.value.message
|
||||
|
||||
def test_error_event_type_with_non_dict_parsed(self) -> None:
|
||||
client = _make_client()
|
||||
with pytest.raises(WorkflowsException) as exc_info:
|
||||
client._parse_sse_data(json.dumps("a plain string"), event_type="error")
|
||||
assert "a plain string" in exc_info.value.message
|
||||
|
||||
def test_malformed_json_raises(self) -> None:
|
||||
client = _make_client()
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
client._parse_sse_data("{not valid json", event_type=None)
|
||||
|
||||
|
||||
class TestIterSSEEvents:
|
||||
@pytest.mark.asyncio
|
||||
async def test_parses_data_lines(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
lines = [f"data: {json.dumps(payload)}"]
|
||||
response = AsyncMock()
|
||||
response.aiter_lines = _async_line_iter(lines)
|
||||
|
||||
events = [e async for e in client._iter_sse_events(response)]
|
||||
assert len(events) == 1
|
||||
assert events[0].stream == "test-stream"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_empty_lines_and_comments(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
lines = ["", ": this is a comment", f"data: {json.dumps(payload)}", ""]
|
||||
response = AsyncMock()
|
||||
response.aiter_lines = _async_line_iter(lines)
|
||||
|
||||
events = [e async for e in client._iter_sse_events(response)]
|
||||
assert len(events) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parses_event_type_and_passes_to_parse(self) -> None:
|
||||
client = _make_client()
|
||||
payload = {"error": "server broke"}
|
||||
lines = ["event: error", f"data: {json.dumps(payload)}"]
|
||||
response = AsyncMock()
|
||||
response.aiter_lines = _async_line_iter(lines)
|
||||
|
||||
with pytest.raises(WorkflowsException) as exc_info:
|
||||
_ = [e async for e in client._iter_sse_events(response)]
|
||||
assert exc_info.value.code == ErrorCode.GET_EVENTS_STREAM_ERROR
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resets_event_type_after_data_line(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
lines = [
|
||||
"event: custom_type",
|
||||
f"data: {json.dumps(payload)}",
|
||||
f"data: {json.dumps(payload)}",
|
||||
]
|
||||
response = AsyncMock()
|
||||
response.aiter_lines = _async_line_iter(lines)
|
||||
|
||||
with patch.object(
|
||||
client, "_parse_sse_data", wraps=client._parse_sse_data
|
||||
) as mock_parse:
|
||||
events = [e async for e in client._iter_sse_events(response)]
|
||||
assert len(events) == 2
|
||||
assert mock_parse.call_args_list[0].args[1] == "custom_type"
|
||||
assert mock_parse.call_args_list[1].args[1] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_non_data_non_event_lines(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
lines = ["id: 123", "retry: 5000", f"data: {json.dumps(payload)}"]
|
||||
response = AsyncMock()
|
||||
response.aiter_lines = _async_line_iter(lines)
|
||||
|
||||
events = [e async for e in client._iter_sse_events(response)]
|
||||
assert len(events) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_failure_logs_warning_and_continues(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
lines = ["data: {not valid json}", f"data: {json.dumps(payload)}"]
|
||||
response = AsyncMock()
|
||||
response.aiter_lines = _async_line_iter(lines)
|
||||
|
||||
with patch("vibe.core.nuage.client.logger") as mock_logger:
|
||||
events = [e async for e in client._iter_sse_events(response)]
|
||||
assert len(events) == 1
|
||||
mock_logger.warning.assert_called_once()
|
||||
|
||||
|
||||
def _setup_mock_client(client: WorkflowsClient, mock_response: AsyncMock) -> None:
|
||||
mock_stream = AsyncMock()
|
||||
mock_stream.__aenter__ = AsyncMock(return_value=mock_response)
|
||||
mock_stream.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_http = AsyncMock()
|
||||
mock_http.stream = lambda *args, **kwargs: mock_stream
|
||||
client._client = mock_http
|
||||
|
||||
|
||||
class TestStreamEvents:
|
||||
@pytest.mark.asyncio
|
||||
async def test_yields_stream_events(self) -> None:
|
||||
client = _make_client()
|
||||
payload = _valid_event_payload()
|
||||
lines = [f"data: {json.dumps(payload)}"]
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.raise_for_status = lambda: None
|
||||
mock_response.aiter_lines = _async_line_iter(lines)
|
||||
_setup_mock_client(client, mock_response)
|
||||
|
||||
params = StreamEventsQueryParams(workflow_exec_id="wf-1", start_seq=0)
|
||||
events = [e async for e in client.stream_events(params)]
|
||||
assert len(events) == 1
|
||||
assert isinstance(events[0], StreamEvent)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reraises_workflows_exception(self) -> None:
|
||||
client = _make_client()
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.raise_for_status = lambda: None
|
||||
mock_response.aiter_lines = _async_line_iter([
|
||||
"event: error",
|
||||
'data: {"error": "stream error"}',
|
||||
])
|
||||
_setup_mock_client(client, mock_response)
|
||||
|
||||
params = StreamEventsQueryParams(workflow_exec_id="wf-1")
|
||||
with pytest.raises(WorkflowsException) as exc_info:
|
||||
_ = [e async for e in client.stream_events(params)]
|
||||
assert exc_info.value.code == ErrorCode.GET_EVENTS_STREAM_ERROR
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wraps_other_exceptions_in_workflows_exception(self) -> None:
|
||||
client = _make_client()
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.raise_for_status.side_effect = RuntimeError("connection lost")
|
||||
_setup_mock_client(client, mock_response)
|
||||
|
||||
params = StreamEventsQueryParams(workflow_exec_id="wf-1")
|
||||
with pytest.raises(WorkflowsException) as exc_info:
|
||||
_ = [e async for e in client.stream_events(params)]
|
||||
assert exc_info.value.code == ErrorCode.GET_EVENTS_STREAM_ERROR
|
||||
assert "Failed to stream events" in exc_info.value.message
|
||||
|
||||
|
||||
def _async_line_iter(lines: list[str]):
|
||||
async def _iter():
|
||||
for line in lines:
|
||||
yield line
|
||||
|
||||
return _iter
|
||||
95
tests/core/test_nuage_client.py
Normal file
95
tests/core/test_nuage_client.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.teleport.nuage import GitHubPublicData, GitHubStatus, NuageClient
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> NuageClient:
|
||||
return NuageClient(
|
||||
base_url="https://test.example.com", api_key="test-key", workflow_id="wf-1"
|
||||
)
|
||||
|
||||
|
||||
class TestWaitForGithubConnection:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_immediately_when_connected(
|
||||
self, client: NuageClient
|
||||
) -> None:
|
||||
connected = GitHubPublicData(status=GitHubStatus.CONNECTED)
|
||||
client.get_github_integration = AsyncMock(return_value=connected)
|
||||
|
||||
result = await client.wait_for_github_connection("exec-1")
|
||||
|
||||
assert result.connected is True
|
||||
client.get_github_integration.assert_called_once_with("exec-1")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_polls_until_connected(self, client: NuageClient) -> None:
|
||||
pending = GitHubPublicData(status=GitHubStatus.PENDING)
|
||||
waiting = GitHubPublicData(
|
||||
status=GitHubStatus.WAITING_FOR_OAUTH, oauth_url="https://github.com/auth"
|
||||
)
|
||||
connected = GitHubPublicData(status=GitHubStatus.CONNECTED)
|
||||
client.get_github_integration = AsyncMock(
|
||||
side_effect=[pending, waiting, connected]
|
||||
)
|
||||
|
||||
with patch("vibe.core.teleport.nuage.asyncio.sleep", new_callable=AsyncMock):
|
||||
result = await client.wait_for_github_connection("exec-1")
|
||||
|
||||
assert result.connected is True
|
||||
assert client.get_github_integration.call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_on_error_status(self, client: NuageClient) -> None:
|
||||
error_data = GitHubPublicData(
|
||||
status=GitHubStatus.ERROR, error="App not installed"
|
||||
)
|
||||
client.get_github_integration = AsyncMock(return_value=error_data)
|
||||
|
||||
with pytest.raises(ServiceTeleportError, match="App not installed"):
|
||||
await client.wait_for_github_connection("exec-1")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_on_oauth_timeout(self, client: NuageClient) -> None:
|
||||
timeout_data = GitHubPublicData(status=GitHubStatus.OAUTH_TIMEOUT)
|
||||
client.get_github_integration = AsyncMock(return_value=timeout_data)
|
||||
|
||||
with pytest.raises(ServiceTeleportError, match="oauth_timeout"):
|
||||
await client.wait_for_github_connection("exec-1")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_on_timeout(self, client: NuageClient) -> None:
|
||||
pending = GitHubPublicData(status=GitHubStatus.PENDING)
|
||||
client.get_github_integration = AsyncMock(return_value=pending)
|
||||
|
||||
monotonic_values = iter([0.0, 0.0, 601.0])
|
||||
with (
|
||||
patch(
|
||||
"vibe.core.teleport.nuage.time.monotonic",
|
||||
side_effect=lambda: next(monotonic_values),
|
||||
),
|
||||
patch("vibe.core.teleport.nuage.asyncio.sleep", new_callable=AsyncMock),
|
||||
pytest.raises(ServiceTeleportError, match="timed out"),
|
||||
):
|
||||
await client.wait_for_github_connection("exec-1", timeout=600.0)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sleeps_with_correct_interval(self, client: NuageClient) -> None:
|
||||
pending = GitHubPublicData(status=GitHubStatus.PENDING)
|
||||
connected = GitHubPublicData(status=GitHubStatus.CONNECTED)
|
||||
client.get_github_integration = AsyncMock(side_effect=[pending, connected])
|
||||
|
||||
with patch(
|
||||
"vibe.core.teleport.nuage.asyncio.sleep", new_callable=AsyncMock
|
||||
) as mock_sleep:
|
||||
await client.wait_for_github_connection("exec-1", interval=5.0)
|
||||
|
||||
mock_sleep.assert_called_once()
|
||||
sleep_duration = mock_sleep.call_args[0][0]
|
||||
assert sleep_duration <= 5.0
|
||||
726
tests/core/test_remote_agent_loop.py
Normal file
726
tests/core/test_remote_agent_loop.py
Normal file
|
|
@ -0,0 +1,726 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from tests.conftest import build_test_vibe_config
|
||||
from vibe.core.nuage.events import (
|
||||
CustomTaskCanceled,
|
||||
CustomTaskCanceledAttributes,
|
||||
CustomTaskCompleted,
|
||||
CustomTaskCompletedAttributes,
|
||||
CustomTaskInProgress,
|
||||
CustomTaskInProgressAttributes,
|
||||
CustomTaskStarted,
|
||||
CustomTaskStartedAttributes,
|
||||
JSONPatchAdd,
|
||||
JSONPatchAppend,
|
||||
JSONPatchPayload,
|
||||
JSONPatchReplace,
|
||||
JSONPayload,
|
||||
)
|
||||
from vibe.core.nuage.remote_events_source import RemoteEventsSource
|
||||
from vibe.core.types import (
|
||||
AssistantEvent,
|
||||
ReasoningEvent,
|
||||
Role,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
ToolStreamEvent,
|
||||
UserMessageEvent,
|
||||
WaitingForInputEvent,
|
||||
)
|
||||
|
||||
_EXEC_ID = "session-123"
|
||||
|
||||
|
||||
def _make_loop(enabled_tools: list[str] | None = None) -> RemoteEventsSource:
|
||||
config = build_test_vibe_config(enabled_tools=enabled_tools or [])
|
||||
return RemoteEventsSource(session_id=_EXEC_ID, config=config)
|
||||
|
||||
|
||||
def _started(
|
||||
task_id: str, task_type: str, payload: dict[str, Any]
|
||||
) -> CustomTaskStarted:
|
||||
return CustomTaskStarted(
|
||||
event_id=f"evt-{task_id}-start",
|
||||
workflow_exec_id=_EXEC_ID,
|
||||
attributes=CustomTaskStartedAttributes(
|
||||
custom_task_id=task_id,
|
||||
custom_task_type=task_type,
|
||||
payload=JSONPayload(value=payload),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _completed(
|
||||
task_id: str, task_type: str, payload: dict[str, Any]
|
||||
) -> CustomTaskCompleted:
|
||||
return CustomTaskCompleted(
|
||||
event_id=f"evt-{task_id}-done",
|
||||
workflow_exec_id=_EXEC_ID,
|
||||
attributes=CustomTaskCompletedAttributes(
|
||||
custom_task_id=task_id,
|
||||
custom_task_type=task_type,
|
||||
payload=JSONPayload(value=payload),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _in_progress(
|
||||
task_id: str, task_type: str, patches: list[Any]
|
||||
) -> CustomTaskInProgress:
|
||||
return CustomTaskInProgress(
|
||||
event_id=f"evt-{task_id}-progress",
|
||||
workflow_exec_id=_EXEC_ID,
|
||||
attributes=CustomTaskInProgressAttributes(
|
||||
custom_task_id=task_id,
|
||||
custom_task_type=task_type,
|
||||
payload=JSONPatchPayload(value=patches),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _canceled(task_id: str, task_type: str, reason: str = "") -> CustomTaskCanceled:
|
||||
return CustomTaskCanceled(
|
||||
event_id=f"evt-{task_id}-cancel",
|
||||
workflow_exec_id=_EXEC_ID,
|
||||
attributes=CustomTaskCanceledAttributes(
|
||||
custom_task_id=task_id, custom_task_type=task_type, reason=reason
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_consume_wait_for_input_event_emits_waiting_event() -> None:
|
||||
loop = _make_loop()
|
||||
event = _started(
|
||||
"wait-task-1",
|
||||
"wait_for_input",
|
||||
{
|
||||
"task_id": "wait-task-1",
|
||||
"input_schema": {"title": "ChatInput"},
|
||||
"label": "What next?",
|
||||
},
|
||||
)
|
||||
|
||||
emitted_events = loop._consume_workflow_event(event)
|
||||
|
||||
assert len(emitted_events) == 2
|
||||
assistant_event = emitted_events[0]
|
||||
waiting_event = emitted_events[1]
|
||||
assert isinstance(assistant_event, AssistantEvent)
|
||||
assert assistant_event.content == "What next?"
|
||||
assert isinstance(waiting_event, WaitingForInputEvent)
|
||||
assert waiting_event.task_id == "wait-task-1"
|
||||
assert waiting_event.label == "What next?"
|
||||
assert waiting_event.predefined_answers is None
|
||||
|
||||
|
||||
def test_consume_agent_input_keeps_repeated_text_across_distinct_turns() -> None:
|
||||
loop = _make_loop()
|
||||
first_event = _completed(
|
||||
"input-1", "AgentInputState", {"input": {"message": [{"text": "continue"}]}}
|
||||
)
|
||||
second_event = _completed(
|
||||
"input-2", "AgentInputState", {"input": {"message": [{"text": "continue"}]}}
|
||||
)
|
||||
|
||||
assert loop._consume_workflow_event(first_event) == []
|
||||
assert loop._consume_workflow_event(second_event) == []
|
||||
|
||||
assert [msg.content for msg in loop.messages if msg.role == Role.user] == [
|
||||
"continue",
|
||||
"continue",
|
||||
]
|
||||
|
||||
|
||||
def test_wait_for_input_emits_predefined_answers_and_user_message() -> None:
|
||||
loop = _make_loop()
|
||||
started = _started(
|
||||
"wait-task-1",
|
||||
"wait_for_input",
|
||||
{
|
||||
"input_schema": {
|
||||
"title": "ChatInput",
|
||||
"properties": {
|
||||
"message": {
|
||||
"examples": [
|
||||
[{"type": "text", "text": "Python"}],
|
||||
[{"type": "text", "text": "JavaScript"}],
|
||||
[{"type": "text", "text": "Other"}],
|
||||
]
|
||||
}
|
||||
},
|
||||
},
|
||||
"label": "Which language?",
|
||||
},
|
||||
)
|
||||
completed = _completed(
|
||||
"wait-task-1",
|
||||
"wait_for_input",
|
||||
{"input": {"message": [{"type": "text", "text": "Python"}]}},
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(started)
|
||||
completed_events = loop._consume_workflow_event(completed)
|
||||
|
||||
assistant_event = next(
|
||||
event for event in started_events if isinstance(event, AssistantEvent)
|
||||
)
|
||||
waiting_event = next(
|
||||
event for event in started_events if isinstance(event, WaitingForInputEvent)
|
||||
)
|
||||
assert assistant_event.content == "Which language?"
|
||||
assert waiting_event.predefined_answers == ["Python", "JavaScript"]
|
||||
user_event = next(
|
||||
event for event in completed_events if isinstance(event, UserMessageEvent)
|
||||
)
|
||||
assert user_event.content == "Python"
|
||||
|
||||
|
||||
def test_tool_events_update_stats_and_messages() -> None:
|
||||
loop = _make_loop(enabled_tools=["todo"])
|
||||
started = _started(
|
||||
"tool-task-1",
|
||||
"AgentToolCallState",
|
||||
{"name": "todo", "tool_call_id": "call-1", "kwargs": {"action": "read"}},
|
||||
)
|
||||
completed = _completed(
|
||||
"tool-task-1",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "todo",
|
||||
"tool_call_id": "call-1",
|
||||
"kwargs": {"action": "read"},
|
||||
"output": {"total_count": 0},
|
||||
},
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(started)
|
||||
completed_events = loop._consume_workflow_event(completed)
|
||||
|
||||
assert any(isinstance(event, ToolCallEvent) for event in started_events)
|
||||
result_event = next(
|
||||
event for event in completed_events if isinstance(event, ToolResultEvent)
|
||||
)
|
||||
assert result_event.error is None
|
||||
assert result_event.cancelled is False
|
||||
assert result_event.tool_call_id == "call-1"
|
||||
assert loop.stats.tool_calls_agreed == 1
|
||||
assert loop.stats.tool_calls_succeeded == 1
|
||||
assert loop.stats.tool_calls_failed == 0
|
||||
tool_messages = [msg for msg in loop.messages if msg.role == Role.tool]
|
||||
assert len(tool_messages) == 1
|
||||
assert tool_messages[0].tool_call_id == "call-1"
|
||||
|
||||
|
||||
def test_ask_user_question_tool_emits_assistant_question() -> None:
|
||||
loop = _make_loop(enabled_tools=["ask_user_question"])
|
||||
started = _started(
|
||||
"tool-task-question",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "ask_user_question",
|
||||
"tool_call_id": "call-question",
|
||||
"kwargs": {
|
||||
"questions": [
|
||||
{
|
||||
"question": "Which file type should I create?",
|
||||
"options": [{"label": "Python"}, {"label": "JavaScript"}],
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
events = loop._consume_workflow_event(started)
|
||||
|
||||
assistant_event = next(
|
||||
event for event in events if isinstance(event, AssistantEvent)
|
||||
)
|
||||
tool_call_event = next(
|
||||
event for event in events if isinstance(event, ToolCallEvent)
|
||||
)
|
||||
assert assistant_event.content == "Which file type should I create?"
|
||||
assert tool_call_event.tool_call_id == "call-question"
|
||||
|
||||
|
||||
def test_ask_user_question_wait_for_input_completion_emits_tool_result() -> None:
|
||||
loop = _make_loop(enabled_tools=["ask_user_question"])
|
||||
ask_started = _started(
|
||||
"tool-task-question",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "ask_user_question",
|
||||
"tool_call_id": "call-question",
|
||||
"kwargs": {
|
||||
"questions": [{"question": "Which type of file?", "options": []}]
|
||||
},
|
||||
},
|
||||
)
|
||||
wait_started = _started(
|
||||
"wait-task-1",
|
||||
"wait_for_input",
|
||||
{"input_schema": {"title": "ChatInput"}, "label": "Which type of file?"},
|
||||
)
|
||||
wait_completed = _completed(
|
||||
"wait-task-1",
|
||||
"wait_for_input",
|
||||
{
|
||||
"input_schema": {"title": "ChatInput"},
|
||||
"label": "Which type of file?",
|
||||
"input": {"message": [{"type": "text", "text": "Python"}]},
|
||||
},
|
||||
)
|
||||
|
||||
loop._consume_workflow_event(ask_started)
|
||||
loop._consume_workflow_event(wait_started)
|
||||
completed_events = loop._consume_workflow_event(wait_completed)
|
||||
|
||||
tool_result = next(
|
||||
(e for e in completed_events if isinstance(e, ToolResultEvent)), None
|
||||
)
|
||||
assert tool_result is not None
|
||||
assert tool_result.tool_call_id == "call-question"
|
||||
user_message = next(e for e in completed_events if isinstance(e, UserMessageEvent))
|
||||
assert user_message.content == "Python"
|
||||
|
||||
|
||||
def test_working_events_without_tool_call_id_render_remote_progress_row() -> None:
|
||||
loop = _make_loop()
|
||||
started = _started(
|
||||
"working-1",
|
||||
"working",
|
||||
{"title": "Creating sandbox", "content": "initializing", "toolUIState": None},
|
||||
)
|
||||
completed = _completed(
|
||||
"working-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Creating sandbox",
|
||||
"content": "sandbox created",
|
||||
"toolUIState": None,
|
||||
},
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(started)
|
||||
completed_events = loop._consume_workflow_event(completed)
|
||||
|
||||
assert started_events == []
|
||||
assert any(isinstance(event, ToolCallEvent) for event in completed_events)
|
||||
assert any(isinstance(event, ToolStreamEvent) for event in completed_events)
|
||||
result_event = next(
|
||||
event for event in completed_events if isinstance(event, ToolResultEvent)
|
||||
)
|
||||
assert result_event.tool_name == "Creating sandbox"
|
||||
assert result_event.tool_call_id == "working-1"
|
||||
|
||||
|
||||
def test_working_events_with_tool_call_id_wait_for_real_tool_call() -> None:
|
||||
loop = _make_loop()
|
||||
working_started = _started(
|
||||
"working-tool-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "writing file",
|
||||
"toolUIState": {"toolCallId": "call-write"},
|
||||
},
|
||||
)
|
||||
tool_started = _started(
|
||||
"tool-task-1",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "write_file",
|
||||
"tool_call_id": "call-write",
|
||||
"kwargs": {
|
||||
"path": "hello_world.js",
|
||||
"content": "console.log('Hello, World!');",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
working_events = loop._consume_workflow_event(working_started)
|
||||
tool_events = loop._consume_workflow_event(tool_started)
|
||||
|
||||
assert any(isinstance(event, ToolCallEvent) for event in working_events)
|
||||
assert any(isinstance(event, ToolStreamEvent) for event in working_events)
|
||||
assert not any(isinstance(event, ToolCallEvent) for event in tool_events)
|
||||
assert not any(isinstance(event, ToolResultEvent) for event in tool_events)
|
||||
|
||||
|
||||
def test_working_task_promoted_to_real_tool_call_does_not_create_duplicate_row() -> (
|
||||
None
|
||||
):
|
||||
loop = _make_loop(enabled_tools=["write_file"])
|
||||
working_started = _started(
|
||||
"working-tool-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Writing file",
|
||||
"content": '# hello.py\n\nprint("Hello, World!")',
|
||||
"toolUIState": None,
|
||||
},
|
||||
)
|
||||
working_promoted = _completed(
|
||||
"working-tool-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "",
|
||||
"toolUIState": {
|
||||
"type": "file",
|
||||
"toolCallId": "call-write",
|
||||
"operations": [
|
||||
{
|
||||
"type": "create",
|
||||
"uri": "/workspace/hello.py",
|
||||
"content": 'print("Hello, World!")',
|
||||
}
|
||||
],
|
||||
},
|
||||
},
|
||||
)
|
||||
agent_tool_completed = _completed(
|
||||
"tool-task-1",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "write_file",
|
||||
"tool_call_id": "call-write",
|
||||
"kwargs": {"path": "hello.py", "content": 'print("Hello, World!")'},
|
||||
"output": {
|
||||
"path": "/workspace/hello.py",
|
||||
"bytes_written": 22,
|
||||
"file_existed": False,
|
||||
"content": 'print("Hello, World!")',
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert loop._consume_workflow_event(working_started) == []
|
||||
|
||||
promoted_events = loop._consume_workflow_event(working_promoted)
|
||||
assert len([e for e in promoted_events if isinstance(e, ToolCallEvent)]) == 1
|
||||
assert not any(isinstance(e, ToolStreamEvent) for e in promoted_events)
|
||||
assert any(isinstance(e, ToolResultEvent) for e in promoted_events)
|
||||
|
||||
completed_events = loop._consume_workflow_event(agent_tool_completed)
|
||||
assert not any(isinstance(e, ToolCallEvent) for e in completed_events)
|
||||
assert not any(isinstance(e, ToolResultEvent) for e in completed_events)
|
||||
|
||||
|
||||
def test_idle_boundary_waits_for_open_tool_results() -> None:
|
||||
loop = _make_loop(enabled_tools=["write_file"])
|
||||
working_started = _started(
|
||||
"working-tool-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "writing file",
|
||||
"toolUIState": {"toolCallId": "call-write"},
|
||||
},
|
||||
)
|
||||
idle_candidate = _completed("input-task-1", "AgentInputState", {"input": None})
|
||||
tool_completed = _completed(
|
||||
"tool-task-1",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "write_file",
|
||||
"tool_call_id": "call-write",
|
||||
"kwargs": {
|
||||
"path": "hello_world.js",
|
||||
"content": "console.log('Hello, World!');",
|
||||
},
|
||||
"output": {
|
||||
"path": "/workspace/hello_world.js",
|
||||
"bytes_written": 29,
|
||||
"file_existed": False,
|
||||
"content": "console.log('Hello, World!');",
|
||||
},
|
||||
},
|
||||
)
|
||||
idle_after_tool = _completed("input-task-2", "AgentInputState", {"input": None})
|
||||
|
||||
working_events = loop._consume_workflow_event(working_started)
|
||||
assert any(isinstance(event, ToolCallEvent) for event in working_events)
|
||||
|
||||
loop._consume_workflow_event(idle_candidate)
|
||||
assert loop._is_idle_boundary(idle_candidate) is False
|
||||
|
||||
tool_events = loop._consume_workflow_event(tool_completed)
|
||||
assert not any(isinstance(event, ToolCallEvent) for event in tool_events)
|
||||
assert any(isinstance(event, ToolResultEvent) for event in tool_events)
|
||||
|
||||
loop._consume_workflow_event(idle_after_tool)
|
||||
assert loop._is_idle_boundary(idle_after_tool) is True
|
||||
|
||||
|
||||
def test_send_user_message_tool_is_not_rendered() -> None:
|
||||
loop = _make_loop()
|
||||
started = _started(
|
||||
"tool-task-send-user-message",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "send_user_message",
|
||||
"tool_call_id": "call-send",
|
||||
"kwargs": {"message": "hello"},
|
||||
},
|
||||
)
|
||||
completed = _completed(
|
||||
"tool-task-send-user-message",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "send_user_message",
|
||||
"tool_call_id": "call-send",
|
||||
"kwargs": {"message": "hello"},
|
||||
"output": {"success": True, "error": None},
|
||||
},
|
||||
)
|
||||
|
||||
assert loop._consume_workflow_event(started) == []
|
||||
assert loop._consume_workflow_event(completed) == []
|
||||
|
||||
|
||||
def test_send_user_message_working_events_are_not_rendered() -> None:
|
||||
loop = _make_loop()
|
||||
started = _started(
|
||||
"working-send-user-message",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing send_user_message",
|
||||
"content": "Hello!",
|
||||
"toolUIState": {"toolCallId": "call-send-working"},
|
||||
},
|
||||
)
|
||||
completed = _completed(
|
||||
"working-send-user-message",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing send_user_message",
|
||||
"content": "Hello!",
|
||||
"toolUIState": {"toolCallId": "call-send-working"},
|
||||
},
|
||||
)
|
||||
|
||||
assert loop._consume_workflow_event(started) == []
|
||||
assert loop._consume_workflow_event(completed) == []
|
||||
|
||||
|
||||
def test_remote_bash_uses_known_tool_display_even_when_disabled_locally() -> None:
|
||||
loop = _make_loop(enabled_tools=["write_file"])
|
||||
started = _started(
|
||||
"tool-task-bash",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "bash",
|
||||
"tool_call_id": "call-bash",
|
||||
"kwargs": {"command": "cat hello.py | wc -c"},
|
||||
},
|
||||
)
|
||||
completed = _completed(
|
||||
"tool-task-bash",
|
||||
"AgentToolCallState",
|
||||
{
|
||||
"name": "bash",
|
||||
"tool_call_id": "call-bash",
|
||||
"kwargs": {"command": "cat hello.py | wc -c"},
|
||||
"output": {
|
||||
"command": "cat hello.py | wc -c",
|
||||
"stdout": "22\n",
|
||||
"stderr": "",
|
||||
"returncode": 0,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(started)
|
||||
completed_events = loop._consume_workflow_event(completed)
|
||||
|
||||
tool_call_event = next(
|
||||
event for event in started_events if isinstance(event, ToolCallEvent)
|
||||
)
|
||||
result_event = next(
|
||||
event for event in completed_events if isinstance(event, ToolResultEvent)
|
||||
)
|
||||
|
||||
assert tool_call_event.tool_name == "bash"
|
||||
assert tool_call_event.tool_class.get_name() == "bash"
|
||||
assert tool_call_event.args is not None
|
||||
assert tool_call_event.args.command == "cat hello.py | wc -c" # type: ignore[attr-defined]
|
||||
assert result_event.result is not None
|
||||
assert result_event.result.command == "cat hello.py | wc -c" # type: ignore[attr-defined]
|
||||
assert result_event.result.stdout == "22\n" # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_canceled_tool_marks_cancelled_and_failed_stats() -> None:
|
||||
loop = _make_loop(enabled_tools=["todo"])
|
||||
loop._task_state["tool-task-2"] = {
|
||||
"name": "todo",
|
||||
"tool_call_id": "call-2",
|
||||
"kwargs": {"action": "read"},
|
||||
}
|
||||
canceled = _canceled(
|
||||
"tool-task-2", "AgentToolCallState", reason="user interrupted tool"
|
||||
)
|
||||
|
||||
events = loop._consume_workflow_event(canceled)
|
||||
|
||||
result_event = next(event for event in events if isinstance(event, ToolResultEvent))
|
||||
assert result_event.cancelled is True
|
||||
assert result_event.error == "Canceled: user interrupted tool"
|
||||
assert loop.stats.tool_calls_failed == 1
|
||||
assert loop.stats.tool_calls_succeeded == 0
|
||||
|
||||
|
||||
def test_working_thinking_type_emits_assistant_events() -> None:
|
||||
loop = _make_loop()
|
||||
started = _started(
|
||||
"thinking-1",
|
||||
"working",
|
||||
{"type": "thinking", "title": "Thinking", "content": "", "toolUIState": None},
|
||||
)
|
||||
in_progress = _in_progress(
|
||||
"thinking-1", "working", [JSONPatchAppend(path="/content", value="Hello!")]
|
||||
)
|
||||
completed = _completed(
|
||||
"thinking-1",
|
||||
"working",
|
||||
{
|
||||
"type": "thinking",
|
||||
"title": "Thinking",
|
||||
"content": "Hello!",
|
||||
"toolUIState": None,
|
||||
},
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(started)
|
||||
progress_events = loop._consume_workflow_event(in_progress)
|
||||
completed_events = loop._consume_workflow_event(completed)
|
||||
|
||||
assert started_events == []
|
||||
assert len(progress_events) == 1
|
||||
assert isinstance(progress_events[0], ReasoningEvent)
|
||||
assert progress_events[0].content == "Hello!"
|
||||
assert completed_events == []
|
||||
|
||||
|
||||
def test_working_bash_progress_without_tool_call_id_streams_command_output() -> None:
|
||||
loop = _make_loop(enabled_tools=["write_file"])
|
||||
started = _started(
|
||||
"working-bash-1",
|
||||
"working",
|
||||
{"type": "tool", "title": "Planning", "content": "", "toolUIState": None},
|
||||
)
|
||||
in_progress = _in_progress(
|
||||
"working-bash-1",
|
||||
"working",
|
||||
[
|
||||
JSONPatchAdd(
|
||||
path="/toolUIState",
|
||||
value={
|
||||
"type": "command",
|
||||
"command": "ls -la /workspace",
|
||||
"result": {
|
||||
"status": "success",
|
||||
"output": "total 4\ndrwxrwxrwx 2 root root 4096 Mar 20 10:18 .\ndrwxr-xr-x 1 root root 80 Mar 20 10:18 ..\n",
|
||||
},
|
||||
},
|
||||
),
|
||||
JSONPatchReplace(path="/title", value="Executing bash"),
|
||||
JSONPatchReplace(path="/content", value=""),
|
||||
],
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(started)
|
||||
progress_events = loop._consume_workflow_event(in_progress)
|
||||
|
||||
assert started_events == []
|
||||
tool_call_event = next(
|
||||
event for event in progress_events if isinstance(event, ToolCallEvent)
|
||||
)
|
||||
tool_stream_event = next(
|
||||
event for event in progress_events if isinstance(event, ToolStreamEvent)
|
||||
)
|
||||
|
||||
assert tool_call_event.tool_name == "bash"
|
||||
assert tool_call_event.tool_class.get_name() == "bash"
|
||||
assert tool_call_event.tool_call_id == "working-bash-1"
|
||||
assert tool_stream_event.tool_name == "bash"
|
||||
assert tool_stream_event.tool_call_id == "working-bash-1"
|
||||
assert "command: ls -la /workspace" in tool_stream_event.message
|
||||
assert "total 4" in tool_stream_event.message
|
||||
assert "drwxrwxrwx 2 root root 4096" in tool_stream_event.message
|
||||
|
||||
|
||||
def test_working_completed_with_tool_call_id_emits_tool_result() -> None:
|
||||
loop = _make_loop(enabled_tools=["write_file"])
|
||||
working_started = _started(
|
||||
"working-tool-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "",
|
||||
"toolUIState": {"toolCallId": "call-write-solo"},
|
||||
},
|
||||
)
|
||||
working_completed = _completed(
|
||||
"working-tool-1",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "",
|
||||
"toolUIState": {
|
||||
"type": "file",
|
||||
"toolCallId": "call-write-solo",
|
||||
"operations": [
|
||||
{
|
||||
"type": "create",
|
||||
"uri": "/workspace/hello.py",
|
||||
"content": 'print("Hello, World!")',
|
||||
}
|
||||
],
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
started_events = loop._consume_workflow_event(working_started)
|
||||
assert any(isinstance(e, ToolCallEvent) for e in started_events)
|
||||
|
||||
completed_events = loop._consume_workflow_event(working_completed)
|
||||
result_events = [e for e in completed_events if isinstance(e, ToolResultEvent)]
|
||||
assert len(result_events) == 1
|
||||
assert result_events[0].error is None
|
||||
assert result_events[0].tool_call_id == "call-write-solo"
|
||||
|
||||
|
||||
def test_working_completed_with_tool_call_id_emits_error_result() -> None:
|
||||
loop = _make_loop(enabled_tools=["write_file"])
|
||||
working_started = _started(
|
||||
"working-tool-2",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "",
|
||||
"toolUIState": {"toolCallId": "call-write-err"},
|
||||
},
|
||||
)
|
||||
working_completed = _completed(
|
||||
"working-tool-2",
|
||||
"working",
|
||||
{
|
||||
"title": "Executing write_file",
|
||||
"content": "Error: File exists. Set overwrite=True.",
|
||||
"toolUIState": {
|
||||
"type": "file",
|
||||
"toolCallId": "call-write-err",
|
||||
"operations": [],
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
loop._consume_workflow_event(working_started)
|
||||
completed_events = loop._consume_workflow_event(working_completed)
|
||||
|
||||
result_events = [e for e in completed_events if isinstance(e, ToolResultEvent)]
|
||||
assert len(result_events) == 1
|
||||
assert result_events[0].error is not None
|
||||
assert result_events[0].tool_call_id == "call-write-err"
|
||||
68
tests/core/test_teleport_get_last_user_message.py
Normal file
68
tests/core/test_teleport_get_last_user_message.py
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.teleport.nuage import TeleportSession
|
||||
from vibe.core.teleport.teleport import TeleportService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def teleport_service() -> TeleportService:
|
||||
return TeleportService.__new__(TeleportService)
|
||||
|
||||
|
||||
def test_returns_last_user_message(teleport_service: TeleportService) -> None:
|
||||
session = TeleportSession(
|
||||
messages=[
|
||||
{"role": "user", "content": "first"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "user", "content": "second"},
|
||||
]
|
||||
)
|
||||
assert teleport_service._get_last_user_message(session) == "second"
|
||||
|
||||
|
||||
def test_returns_none_when_no_user_messages(teleport_service: TeleportService) -> None:
|
||||
session = TeleportSession(
|
||||
messages=[
|
||||
{"role": "system", "content": "system prompt"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
]
|
||||
)
|
||||
assert teleport_service._get_last_user_message(session) is None
|
||||
|
||||
|
||||
def test_returns_none_for_empty_messages(teleport_service: TeleportService) -> None:
|
||||
session = TeleportSession(messages=[])
|
||||
assert teleport_service._get_last_user_message(session) is None
|
||||
|
||||
|
||||
def test_skips_non_string_content(teleport_service: TeleportService) -> None:
|
||||
session = TeleportSession(
|
||||
messages=[
|
||||
{"role": "user", "content": [{"type": "text", "text": "block content"}]}
|
||||
]
|
||||
)
|
||||
assert teleport_service._get_last_user_message(session) is None
|
||||
|
||||
|
||||
def test_skips_empty_string_content(teleport_service: TeleportService) -> None:
|
||||
session = TeleportSession(messages=[{"role": "user", "content": ""}])
|
||||
assert teleport_service._get_last_user_message(session) is None
|
||||
|
||||
|
||||
def test_skips_non_string_returns_earlier_string(
|
||||
teleport_service: TeleportService,
|
||||
) -> None:
|
||||
session = TeleportSession(
|
||||
messages=[
|
||||
{"role": "user", "content": "valid message"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "block"}]},
|
||||
]
|
||||
)
|
||||
assert teleport_service._get_last_user_message(session) == "valid message"
|
||||
|
||||
|
||||
def test_skips_missing_content_key(teleport_service: TeleportService) -> None:
|
||||
session = TeleportSession(messages=[{"role": "user"}])
|
||||
assert teleport_service._get_last_user_message(session) is None
|
||||
|
|
@ -5,73 +5,73 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import httpx
|
||||
import pytest
|
||||
|
||||
from vibe.core.auth import EncryptedPayload
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.teleport.nuage import (
|
||||
CreateLeChatThreadInput,
|
||||
GitRepoConfig,
|
||||
ChatAssistantParams,
|
||||
GitHubParams,
|
||||
NuageClient,
|
||||
VibeNewSandbox,
|
||||
VibeSandboxConfig,
|
||||
WorkflowIntegrations,
|
||||
WorkflowParams,
|
||||
)
|
||||
|
||||
|
||||
class TestNuageModels:
|
||||
def test_git_repo_config_defaults(self) -> None:
|
||||
config = GitRepoConfig(url="https://github.com/owner/repo.git")
|
||||
assert config.url == "https://github.com/owner/repo.git"
|
||||
assert config.branch is None
|
||||
assert config.commit is None
|
||||
def test_github_params_defaults(self) -> None:
|
||||
params = GitHubParams(repo="owner/repo")
|
||||
assert params.repo == "owner/repo"
|
||||
assert params.branch is None
|
||||
assert params.commit is None
|
||||
assert params.pr_number is None
|
||||
assert params.teleported_diffs is None
|
||||
|
||||
def test_git_repo_config_with_values(self) -> None:
|
||||
config = GitRepoConfig(
|
||||
url="https://github.com/owner/repo.git", branch="main", commit="abc123"
|
||||
def test_github_params_with_values(self) -> None:
|
||||
params = GitHubParams(
|
||||
repo="owner/repo",
|
||||
branch="main",
|
||||
commit="abc123",
|
||||
pr_number=42,
|
||||
teleported_diffs=b"base64data",
|
||||
)
|
||||
assert config.branch == "main"
|
||||
assert config.commit == "abc123"
|
||||
assert params.repo == "owner/repo"
|
||||
assert params.branch == "main"
|
||||
assert params.commit == "abc123"
|
||||
assert params.pr_number == 42
|
||||
assert params.teleported_diffs == b"base64data"
|
||||
|
||||
def test_vibe_sandbox_config_defaults(self) -> None:
|
||||
config = VibeSandboxConfig()
|
||||
assert config.git_repo is None
|
||||
def test_chat_assistant_params(self) -> None:
|
||||
params = ChatAssistantParams(
|
||||
user_message="test message", project_name="test-project"
|
||||
)
|
||||
assert params.user_message == "test message"
|
||||
assert params.project_name == "test-project"
|
||||
|
||||
def test_vibe_new_sandbox_defaults(self) -> None:
|
||||
sandbox = VibeNewSandbox()
|
||||
assert sandbox.type == "new"
|
||||
assert sandbox.config.git_repo is None
|
||||
assert sandbox.teleported_diffs is None
|
||||
def test_workflow_integrations(self) -> None:
|
||||
integrations = WorkflowIntegrations(
|
||||
github=GitHubParams(repo="owner/repo"),
|
||||
chat_assistant=ChatAssistantParams(user_message="test"),
|
||||
)
|
||||
assert integrations.github is not None
|
||||
assert integrations.chat_assistant is not None
|
||||
|
||||
def test_workflow_params_serialization(self) -> None:
|
||||
params = WorkflowParams(
|
||||
prompt="test prompt",
|
||||
sandbox=VibeNewSandbox(
|
||||
config=VibeSandboxConfig(
|
||||
git_repo=GitRepoConfig(
|
||||
url="https://github.com/owner/repo.git",
|
||||
branch="main",
|
||||
commit="abc123",
|
||||
)
|
||||
integrations=WorkflowIntegrations(
|
||||
github=GitHubParams(
|
||||
repo="owner/repo",
|
||||
branch="main",
|
||||
commit="abc123",
|
||||
pr_number=42,
|
||||
teleported_diffs=b"base64data",
|
||||
),
|
||||
teleported_diffs=b"base64data",
|
||||
chat_assistant=ChatAssistantParams(user_message="test"),
|
||||
),
|
||||
)
|
||||
data = params.model_dump()
|
||||
assert data["prompt"] == "test prompt"
|
||||
assert data["sandbox"]["type"] == "new"
|
||||
assert data["sandbox"]["config"]["git_repo"]["url"] == (
|
||||
"https://github.com/owner/repo.git"
|
||||
)
|
||||
assert data["sandbox"]["teleported_diffs"] == b"base64data"
|
||||
|
||||
def test_create_le_chat_thread_input(self) -> None:
|
||||
input_data = CreateLeChatThreadInput(
|
||||
encrypted_api_key={"key": "value"},
|
||||
user_message="test message",
|
||||
project_name="test-project",
|
||||
)
|
||||
assert input_data.encrypted_api_key == {"key": "value"}
|
||||
assert input_data.user_message == "test message"
|
||||
assert input_data.project_name == "test-project"
|
||||
assert data["integrations"]["github"]["repo"] == "owner/repo"
|
||||
assert data["integrations"]["github"]["pr_number"] == 42
|
||||
assert data["integrations"]["github"]["teleported_diffs"] == b"base64data"
|
||||
|
||||
|
||||
class TestNuageClientContextManager:
|
||||
|
|
@ -116,7 +116,7 @@ class TestNuageClientStartWorkflow:
|
|||
mock_response.json.return_value = {"execution_id": "exec-123"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
params = WorkflowParams(prompt="test", sandbox=VibeNewSandbox())
|
||||
params = WorkflowParams(prompt="test", integrations=WorkflowIntegrations())
|
||||
execution_id = await nuage.start_workflow(params)
|
||||
|
||||
assert execution_id == "exec-123"
|
||||
|
|
@ -133,25 +133,12 @@ class TestNuageClientStartWorkflow:
|
|||
mock_response.text = "Internal Server Error"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
params = WorkflowParams(prompt="test", sandbox=VibeNewSandbox())
|
||||
params = WorkflowParams(prompt="test", integrations=WorkflowIntegrations())
|
||||
with pytest.raises(ServiceTeleportError, match="Nuage workflow trigger failed"):
|
||||
await nuage.start_workflow(params)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_workflow_unauthorized_hint(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Unauthorized"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
params = WorkflowParams(prompt="test", sandbox=VibeNewSandbox())
|
||||
with pytest.raises(ServiceTeleportError, match="STAGING_MISTRAL_API_KEY"):
|
||||
await nuage.start_workflow(params)
|
||||
|
||||
|
||||
class TestNuageClientSendGithubToken:
|
||||
class TestNuageClientGetGitHubIntegration:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
|
@ -163,48 +150,43 @@ class TestNuageClientSendGithubToken:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_github_token_success(
|
||||
self,
|
||||
nuage: NuageClient,
|
||||
mock_client: MagicMock,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
async def test_get_github_integration_connected(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
query_response = MagicMock()
|
||||
query_response.is_success = True
|
||||
query_response.json.return_value = {
|
||||
"result": {
|
||||
"public_key": (
|
||||
"-----BEGIN PUBLIC KEY-----\n"
|
||||
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA0Z3VS5JJcds3xfn/"
|
||||
"ygWyf8TFXQNZ0XsLOqXB1Mi2+bKPFv1WfhECTxJ3c6SXl0p1sGyWTFxRV8u/"
|
||||
"bKqYZ0E6VZ6YRTRPFiGq0kkONjVBFxOQ8Y0jeT0d9e0Y3E3MWDL8tQ0Nz9v8"
|
||||
"5Y7gC8F1m/dEbBwPjCJQV0Dg0z3gZDO8RCG0GrBoLO0b+NNqL8FXPPDXQ1l4"
|
||||
"FGnYM0gZ1rCU7Y/zTN1wI4sCQ0GJQPDA1hWB8KRJl5x0ZDXE3rRwT1E8c+Fn"
|
||||
"ZFV1nN0C6zxF7GpVY3FVWXS4PA0FH+8C1+TnYgBL7xS0o+LF6PgjGT5F3CXD"
|
||||
"BZmYSxKL+EsVVGT5EuYbJE9TxVwIDAQAB\n"
|
||||
"-----END PUBLIC KEY-----"
|
||||
)
|
||||
}
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = True
|
||||
mock_response.json.return_value = {
|
||||
"result": {"status": "connected", "oauth_url": None, "error": None}
|
||||
}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
signal_response = MagicMock()
|
||||
signal_response.is_success = True
|
||||
result = await nuage.get_github_integration("exec-123")
|
||||
|
||||
mock_client.post = AsyncMock(side_effect=[query_response, signal_response])
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.teleport.nuage.encrypt", lambda _token, _key: encrypted
|
||||
)
|
||||
|
||||
await nuage.send_github_token("exec-123", "ghp_token")
|
||||
|
||||
assert mock_client.post.call_count == 2
|
||||
assert result.connected is True
|
||||
assert result.oauth_url is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_public_key_failure(
|
||||
async def test_get_github_integration_waiting_for_oauth(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = True
|
||||
mock_response.json.return_value = {
|
||||
"result": {
|
||||
"status": "waiting_for_oauth",
|
||||
"oauth_url": "https://github.com/login/oauth",
|
||||
"error": None,
|
||||
}
|
||||
}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
result = await nuage.get_github_integration("exec-123")
|
||||
|
||||
assert result.connected is False
|
||||
assert result.oauth_url == "https://github.com/login/oauth"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_github_integration_failure(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -212,26 +194,13 @@ class TestNuageClientSendGithubToken:
|
|||
mock_response.text = "Not found"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with pytest.raises(ServiceTeleportError, match="Failed to get public key"):
|
||||
await nuage._query_public_key("exec-123")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_signal_encrypted_token_failure(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Signal failed"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
with pytest.raises(ServiceTeleportError, match="Failed to send GitHub token"):
|
||||
await nuage._signal_encrypted_token("exec-123", encrypted)
|
||||
with pytest.raises(
|
||||
ServiceTeleportError, match="Failed to get GitHub integration"
|
||||
):
|
||||
await nuage.get_github_integration("exec-123")
|
||||
|
||||
|
||||
class TestNuageClientCreateLeChatThread:
|
||||
class TestNuageClientGetChatAssistantUrl:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
|
@ -243,90 +212,32 @@ class TestNuageClientCreateLeChatThread:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_le_chat_thread_success(
|
||||
self,
|
||||
nuage: NuageClient,
|
||||
mock_client: MagicMock,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
async def test_get_chat_assistant_url_success(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
query_response = MagicMock()
|
||||
query_response.is_success = True
|
||||
query_response.json.return_value = {
|
||||
"result": {
|
||||
"public_key": (
|
||||
"-----BEGIN PUBLIC KEY-----\n"
|
||||
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA0Z3VS5JJcds3xfn/"
|
||||
"ygWyf8TFXQNZ0XsLOqXB1Mi2+bKPFv1WfhECTxJ3c6SXl0p1sGyWTFxRV8u/"
|
||||
"bKqYZ0E6VZ6YRTRPFiGq0kkONjVBFxOQ8Y0jeT0d9e0Y3E3MWDL8tQ0Nz9v8"
|
||||
"5Y7gC8F1m/dEbBwPjCJQV0Dg0z3gZDO8RCG0GrBoLO0b+NNqL8FXPPDXQ1l4"
|
||||
"FGnYM0gZ1rCU7Y/zTN1wI4sCQ0GJQPDA1hWB8KRJl5x0ZDXE3rRwT1E8c+Fn"
|
||||
"ZFV1nN0C6zxF7GpVY3FVWXS4PA0FH+8C1+TnYgBL7xS0o+LF6PgjGT5F3CXD"
|
||||
"BZmYSxKL+EsVVGT5EuYbJE9TxVwIDAQAB\n"
|
||||
"-----END PUBLIC KEY-----"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
update_response = MagicMock()
|
||||
update_response.is_success = True
|
||||
update_response.json.return_value = {
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = True
|
||||
mock_response.json.return_value = {
|
||||
"result": {"chat_url": "https://chat.example.com/thread/123"}
|
||||
}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_client.post = AsyncMock(side_effect=[query_response, update_response])
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.teleport.nuage.encrypt", lambda _token, _key: encrypted
|
||||
)
|
||||
|
||||
url = await nuage.create_le_chat_thread("exec-123", "test message")
|
||||
url = await nuage.get_chat_assistant_url("exec-123")
|
||||
assert url == "https://chat.example.com/thread/123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_le_chat_thread_failure(
|
||||
self,
|
||||
nuage: NuageClient,
|
||||
mock_client: MagicMock,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
async def test_get_chat_assistant_url_failure(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
query_response = MagicMock()
|
||||
query_response.is_success = True
|
||||
query_response.json.return_value = {
|
||||
"result": {
|
||||
"public_key": (
|
||||
"-----BEGIN PUBLIC KEY-----\n"
|
||||
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA0Z3VS5JJcds3xfn/"
|
||||
"ygWyf8TFXQNZ0XsLOqXB1Mi2+bKPFv1WfhECTxJ3c6SXl0p1sGyWTFxRV8u/"
|
||||
"bKqYZ0E6VZ6YRTRPFiGq0kkONjVBFxOQ8Y0jeT0d9e0Y3E3MWDL8tQ0Nz9v8"
|
||||
"5Y7gC8F1m/dEbBwPjCJQV0Dg0z3gZDO8RCG0GrBoLO0b+NNqL8FXPPDXQ1l4"
|
||||
"FGnYM0gZ1rCU7Y/zTN1wI4sCQ0GJQPDA1hWB8KRJl5x0ZDXE3rRwT1E8c+Fn"
|
||||
"ZFV1nN0C6zxF7GpVY3FVWXS4PA0FH+8C1+TnYgBL7xS0o+LF6PgjGT5F3CXD"
|
||||
"BZmYSxKL+EsVVGT5EuYbJE9TxVwIDAQAB\n"
|
||||
"-----END PUBLIC KEY-----"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
update_response = MagicMock()
|
||||
update_response.is_success = False
|
||||
update_response.text = "Failed to create thread"
|
||||
|
||||
mock_client.post = AsyncMock(side_effect=[query_response, update_response])
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.teleport.nuage.encrypt", lambda _token, _key: encrypted
|
||||
)
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Failed"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with pytest.raises(
|
||||
ServiceTeleportError, match="Failed to create Le Chat thread"
|
||||
ServiceTeleportError, match="Failed to get chat assistant integration"
|
||||
):
|
||||
await nuage.create_le_chat_thread("exec-123", "test message")
|
||||
await nuage.get_chat_assistant_url("exec-123")
|
||||
|
||||
|
||||
class TestNuageClientHeaders:
|
||||
|
|
|
|||
|
|
@ -16,18 +16,19 @@ from vibe.core.teleport.errors import (
|
|||
ServiceTeleportNotSupportedError,
|
||||
)
|
||||
from vibe.core.teleport.git import GitRepoInfo
|
||||
from vibe.core.teleport.nuage import TeleportSession
|
||||
from vibe.core.teleport.teleport import _NUAGE_EXECUTION_URL_TEMPLATE, TeleportService
|
||||
from vibe.core.teleport.nuage import GitHubStatus, TeleportSession
|
||||
from vibe.core.teleport.teleport import TeleportService
|
||||
from vibe.core.teleport.types import (
|
||||
TeleportAuthCompleteEvent,
|
||||
TeleportAuthRequiredEvent,
|
||||
TeleportCheckingGitEvent,
|
||||
TeleportCompleteEvent,
|
||||
TeleportFetchingUrlEvent,
|
||||
TeleportPushingEvent,
|
||||
TeleportPushRequiredEvent,
|
||||
TeleportPushResponseEvent,
|
||||
TeleportSendingGithubTokenEvent,
|
||||
TeleportStartingWorkflowEvent,
|
||||
TeleportWaitingForGitHubEvent,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -68,7 +69,7 @@ class TestTeleportServiceCompressDiff:
|
|||
service._compress_diff(large_diff, max_size=100)
|
||||
|
||||
|
||||
class TestTeleportServiceBuildSandbox:
|
||||
class TestTeleportServiceBuildGitHubParams:
|
||||
@pytest.fixture
|
||||
def service(self, tmp_path: Path) -> TeleportService:
|
||||
mock_session_logger = MagicMock()
|
||||
|
|
@ -80,7 +81,7 @@ class TestTeleportServiceBuildSandbox:
|
|||
workdir=tmp_path,
|
||||
)
|
||||
|
||||
def test_builds_sandbox_from_git_info(self, service: TeleportService) -> None:
|
||||
def test_builds_params_from_git_info(self, service: TeleportService) -> None:
|
||||
git_info = GitRepoInfo(
|
||||
remote_url="https://github.com/owner/repo.git",
|
||||
owner="owner",
|
||||
|
|
@ -89,14 +90,12 @@ class TestTeleportServiceBuildSandbox:
|
|||
commit="abc123",
|
||||
diff="",
|
||||
)
|
||||
sandbox = service._build_sandbox(git_info)
|
||||
params = service._build_github_params(git_info)
|
||||
|
||||
assert sandbox.type == "new"
|
||||
assert sandbox.config.git_repo is not None
|
||||
assert sandbox.config.git_repo.url == "https://github.com/owner/repo.git"
|
||||
assert sandbox.config.git_repo.branch == "main"
|
||||
assert sandbox.config.git_repo.commit == "abc123"
|
||||
assert sandbox.teleported_diffs is None
|
||||
assert params.repo == "owner/repo"
|
||||
assert params.branch == "main"
|
||||
assert params.commit == "abc123"
|
||||
assert params.teleported_diffs is None
|
||||
|
||||
def test_includes_compressed_diff(self, service: TeleportService) -> None:
|
||||
git_info = GitRepoInfo(
|
||||
|
|
@ -107,9 +106,9 @@ class TestTeleportServiceBuildSandbox:
|
|||
commit="abc123",
|
||||
diff="diff content",
|
||||
)
|
||||
sandbox = service._build_sandbox(git_info)
|
||||
params = service._build_github_params(git_info)
|
||||
|
||||
assert sandbox.teleported_diffs is not None
|
||||
assert params.teleported_diffs is not None
|
||||
|
||||
|
||||
class TestTeleportServiceValidateConfig:
|
||||
|
|
@ -122,9 +121,7 @@ class TestTeleportServiceValidateConfig:
|
|||
nuage_api_key="",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
with pytest.raises(
|
||||
ServiceTeleportError, match="STAGING_MISTRAL_API_KEY not set"
|
||||
):
|
||||
with pytest.raises(ServiceTeleportError, match="MISTRAL_API_KEY not set"):
|
||||
service._validate_config()
|
||||
|
||||
def test_passes_when_api_key_set(self, tmp_path: Path) -> None:
|
||||
|
|
@ -138,6 +135,21 @@ class TestTeleportServiceValidateConfig:
|
|||
)
|
||||
service._validate_config()
|
||||
|
||||
def test_uses_custom_env_var_name_in_error(self, tmp_path: Path) -> None:
|
||||
mock_session_logger = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
mock_config.nuage_api_key_env_var = "CUSTOM_API_KEY"
|
||||
service = TeleportService(
|
||||
session_logger=mock_session_logger,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_api_key="",
|
||||
workdir=tmp_path,
|
||||
vibe_config=mock_config,
|
||||
)
|
||||
with pytest.raises(ServiceTeleportError, match="CUSTOM_API_KEY not set"):
|
||||
service._validate_config()
|
||||
|
||||
|
||||
class TestTeleportServiceCheckSupported:
|
||||
@pytest.fixture
|
||||
|
|
@ -213,6 +225,8 @@ class TestTeleportServiceExecute:
|
|||
nuage_api_key="api-key",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
service._git.fetch = AsyncMock()
|
||||
service._git.is_branch_pushed = AsyncMock(return_value=True)
|
||||
return service
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -226,20 +240,32 @@ class TestTeleportServiceExecute:
|
|||
diff="",
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_github_connected(self) -> MagicMock:
|
||||
github_data = MagicMock()
|
||||
github_data.connected = True
|
||||
github_data.oauth_url = None
|
||||
github_data.status = GitHubStatus.CONNECTED
|
||||
return github_data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_happy_path_commit_pushed_with_token(
|
||||
self, service: TeleportService, git_info: GitRepoInfo
|
||||
async def test_execute_happy_path_github_already_connected(
|
||||
self,
|
||||
service: TeleportService,
|
||||
git_info: GitRepoInfo,
|
||||
mock_github_connected: MagicMock,
|
||||
) -> None:
|
||||
service._git.get_info = AsyncMock(return_value=git_info)
|
||||
service._git.is_commit_pushed = AsyncMock(return_value=True)
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value="ghp_existing_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
mock_nuage.get_github_integration = AsyncMock(
|
||||
return_value=mock_github_connected
|
||||
)
|
||||
mock_nuage.get_chat_assistant_url = AsyncMock(
|
||||
return_value="https://chat.example.com/123"
|
||||
)
|
||||
service._nuage = mock_nuage
|
||||
|
||||
session = TeleportSession()
|
||||
|
|
@ -250,29 +276,31 @@ class TestTeleportServiceExecute:
|
|||
|
||||
assert isinstance(events[0], TeleportCheckingGitEvent)
|
||||
assert isinstance(events[1], TeleportStartingWorkflowEvent)
|
||||
assert isinstance(events[2], TeleportSendingGithubTokenEvent)
|
||||
assert isinstance(events[3], TeleportCompleteEvent)
|
||||
expected_url = _NUAGE_EXECUTION_URL_TEMPLATE.format(
|
||||
workflow_id="workflow-id", execution_id="exec-123"
|
||||
)
|
||||
assert events[3].url == expected_url
|
||||
assert isinstance(events[2], TeleportWaitingForGitHubEvent)
|
||||
assert isinstance(events[3], TeleportFetchingUrlEvent)
|
||||
assert isinstance(events[4], TeleportCompleteEvent)
|
||||
assert events[4].url == "https://chat.example.com/123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_requires_push_and_user_approves(
|
||||
self, service: TeleportService, git_info: GitRepoInfo
|
||||
self,
|
||||
service: TeleportService,
|
||||
git_info: GitRepoInfo,
|
||||
mock_github_connected: MagicMock,
|
||||
) -> None:
|
||||
service._git.get_info = AsyncMock(return_value=git_info)
|
||||
service._git.is_commit_pushed = AsyncMock(return_value=False)
|
||||
service._git.get_unpushed_commit_count = AsyncMock(return_value=3)
|
||||
service._git.push_current_branch = AsyncMock(return_value=True)
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value="ghp_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
mock_nuage.get_github_integration = AsyncMock(
|
||||
return_value=mock_github_connected
|
||||
)
|
||||
mock_nuage.get_chat_assistant_url = AsyncMock(
|
||||
return_value="https://chat.example.com/123"
|
||||
)
|
||||
service._nuage = mock_nuage
|
||||
|
||||
session = TeleportSession()
|
||||
|
|
@ -315,25 +343,29 @@ class TestTeleportServiceExecute:
|
|||
await gen.asend(TeleportPushResponseEvent(approved=False))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_requires_auth_flow(
|
||||
async def test_execute_requires_oauth_flow(
|
||||
self, service: TeleportService, git_info: GitRepoInfo
|
||||
) -> None:
|
||||
service._git.get_info = AsyncMock(return_value=git_info)
|
||||
service._git.is_commit_pushed = AsyncMock(return_value=True)
|
||||
|
||||
mock_handle = MagicMock()
|
||||
mock_handle.info.user_code = "ABC-123"
|
||||
mock_handle.info.verification_uri = "https://github.com/login/device"
|
||||
github_pending = MagicMock()
|
||||
github_pending.connected = False
|
||||
github_pending.oauth_url = "https://github.com/login/oauth"
|
||||
github_pending.status = GitHubStatus.WAITING_FOR_OAUTH
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value=None)
|
||||
mock_github_auth.start_device_flow = AsyncMock(return_value=mock_handle)
|
||||
mock_github_auth.wait_for_token = AsyncMock(return_value="ghp_new_token")
|
||||
service._github_auth = mock_github_auth
|
||||
github_connected = MagicMock()
|
||||
github_connected.connected = True
|
||||
github_connected.oauth_url = None
|
||||
github_connected.status = GitHubStatus.CONNECTED
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
mock_nuage.get_github_integration = AsyncMock(return_value=github_pending)
|
||||
mock_nuage.wait_for_github_connection = AsyncMock(return_value=github_connected)
|
||||
mock_nuage.get_chat_assistant_url = AsyncMock(
|
||||
return_value="https://chat.example.com/123"
|
||||
)
|
||||
service._nuage = mock_nuage
|
||||
|
||||
session = TeleportSession()
|
||||
|
|
@ -343,35 +375,42 @@ class TestTeleportServiceExecute:
|
|||
events.append(event)
|
||||
|
||||
assert isinstance(events[0], TeleportCheckingGitEvent)
|
||||
assert isinstance(events[1], TeleportAuthRequiredEvent)
|
||||
assert events[1].user_code == "ABC-123"
|
||||
assert isinstance(events[2], TeleportAuthCompleteEvent)
|
||||
assert isinstance(events[3], TeleportStartingWorkflowEvent)
|
||||
assert isinstance(events[1], TeleportStartingWorkflowEvent)
|
||||
assert isinstance(events[2], TeleportWaitingForGitHubEvent)
|
||||
assert isinstance(events[3], TeleportAuthRequiredEvent)
|
||||
assert events[3].oauth_url == "https://github.com/login/oauth"
|
||||
assert isinstance(events[4], TeleportAuthCompleteEvent)
|
||||
assert isinstance(events[-1], TeleportCompleteEvent)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uses_default_prompt_when_none(
|
||||
self, service: TeleportService, git_info: GitRepoInfo
|
||||
self,
|
||||
service: TeleportService,
|
||||
git_info: GitRepoInfo,
|
||||
mock_github_connected: MagicMock,
|
||||
) -> None:
|
||||
service._git.get_info = AsyncMock(return_value=git_info)
|
||||
service._git.is_commit_pushed = AsyncMock(return_value=True)
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value="ghp_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
mock_nuage.get_github_integration = AsyncMock(
|
||||
return_value=mock_github_connected
|
||||
)
|
||||
mock_nuage.get_chat_assistant_url = AsyncMock(
|
||||
return_value="https://chat.example.com/123"
|
||||
)
|
||||
service._nuage = mock_nuage
|
||||
|
||||
session = TeleportSession()
|
||||
session = TeleportSession(
|
||||
messages=[{"role": "user", "content": "help me refactor"}]
|
||||
)
|
||||
gen = service.execute(None, session)
|
||||
async for _ in gen:
|
||||
pass
|
||||
|
||||
call_args = mock_nuage.start_workflow.call_args
|
||||
assert "continue where you left off" in call_args[0][0].prompt
|
||||
assert "teleported" in call_args[0][0].prompt.lower()
|
||||
|
||||
|
||||
class TestTeleportServiceContextManager:
|
||||
|
|
@ -388,7 +427,6 @@ class TestTeleportServiceContextManager:
|
|||
assert service._client is None
|
||||
async with service:
|
||||
assert service._client is not None
|
||||
assert service._github_auth is not None
|
||||
assert service._nuage is not None
|
||||
assert service._client is None
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue