Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai>
Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Sirieix 2026-05-27 21:39:06 +02:00 committed by GitHub
parent cf3f4ca58f
commit 1843196d88
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 399 additions and 1675 deletions

View file

@ -93,7 +93,6 @@ async def test_teleport_command_without_history_sends_early_failure_telemetry(
"stage": "no_history",
"error_class": "TeleportNoHistoryError",
"push_required": False,
"github_auth_required": False,
"nb_session_messages": 0,
"session_id": app.agent_loop.session_id,
},
@ -134,7 +133,6 @@ async def test_teleport_command_in_remote_session_sends_early_failure_telemetry(
"stage": "remote_session",
"error_class": "TeleportRemoteSessionError",
"push_required": False,
"github_auth_required": False,
"nb_session_messages": 2,
"session_id": app.agent_loop.session_id,
},

View file

@ -0,0 +1,45 @@
from __future__ import annotations
from unittest.mock import MagicMock
from vibe.core.agent_loop import AgentLoop
from vibe.core.types import LLMMessage, Role
def _call(*messages: LLMMessage) -> LLMMessage | None:
fake_self = MagicMock(spec=["messages"])
fake_self.messages = list(messages)
return AgentLoop._last_user_message(fake_self)
def test_returns_last_user_message() -> None:
m1 = LLMMessage(role=Role.user, content="first")
m2 = LLMMessage(role=Role.assistant, content="reply")
m3 = LLMMessage(role=Role.user, content="second")
result = _call(m1, m2, m3)
assert result is m3
def test_returns_none_when_no_user_messages() -> None:
assert (
_call(
LLMMessage(role=Role.system, content="system prompt"),
LLMMessage(role=Role.assistant, content="hello"),
)
is None
)
def test_returns_none_for_empty_messages() -> None:
assert _call() is None
def test_skips_injected_user_messages() -> None:
real = LLMMessage(role=Role.user, content="real")
injected = LLMMessage(role=Role.user, content="reminder", injected=True)
assert _call(real, injected) is real
def test_returns_user_message_with_empty_content() -> None:
empty = LLMMessage(role=Role.user, content="")
assert _call(empty) is empty

View file

@ -1,125 +0,0 @@
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)
results = [data async for data in client.wait_for_github_connection("exec-1")]
assert len(results) == 1
assert results[0].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):
results = [
data async for data in client.wait_for_github_connection("exec-1")
]
assert len(results) == 3
assert results[-1].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"):
async for _ in client.wait_for_github_connection("exec-1"):
pass
@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"):
async for _ in client.wait_for_github_connection("exec-1"):
pass
@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"),
):
async for _ in client.wait_for_github_connection("exec-1", timeout=600.0):
pass
@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:
async for _ in client.wait_for_github_connection("exec-1", interval=5.0):
pass
mock_sleep.assert_called_once()
sleep_duration = mock_sleep.call_args[0][0]
assert sleep_duration <= 5.0
@pytest.mark.asyncio
async def test_yields_each_poll_result(self, client: NuageClient) -> None:
pending = GitHubPublicData(
status=GitHubStatus.WAITING_FOR_OAUTH,
error="Please connect GitHub",
oauth_url="https://github.com/auth",
)
connected = GitHubPublicData(status=GitHubStatus.CONNECTED)
client.get_github_integration = AsyncMock(
side_effect=[pending, pending, connected]
)
with patch("vibe.core.teleport.nuage.asyncio.sleep", new_callable=AsyncMock):
results = [
data async for data in client.wait_for_github_connection("exec-1")
]
assert len(results) == 3
assert results[0].error == "Please connect GitHub"
assert results[1].error == "Please connect GitHub"
assert results[2].connected is True

View file

@ -368,15 +368,12 @@ class TestTelemetryClient:
config = build_test_vibe_config(enable_telemetry=True)
client = TelemetryClient(config_getter=lambda: config)
client.send_teleport_completed(
push_required=True, github_auth_required=False, nb_session_messages=4
)
client.send_teleport_completed(push_required=True, nb_session_messages=4)
assert len(telemetry_events) == 1
assert telemetry_events[0]["event_name"] == "vibe.teleport_completed"
assert telemetry_events[0]["properties"] == {
"push_required": True,
"github_auth_required": False,
"nb_session_messages": 4,
}
@ -390,7 +387,6 @@ class TestTelemetryClient:
stage="push",
error_class="ServiceTeleportError",
push_required=True,
github_auth_required=False,
nb_session_messages=4,
)
@ -400,7 +396,6 @@ class TestTelemetryClient:
"stage": "push",
"error_class": "ServiceTeleportError",
"push_required": True,
"github_auth_required": False,
"nb_session_messages": 4,
}

View file

@ -1,209 +0,0 @@
from __future__ import annotations
import json
import httpx
import pytest
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.teleport.experimental_nuage import (
ExperimentalNuageClient,
ExperimentalNuageContext,
ExperimentalNuageDiff,
ExperimentalNuageMessage,
ExperimentalNuageRepository,
ExperimentalNuageRequest,
ExperimentalNuageTextPart,
)
def _request() -> ExperimentalNuageRequest:
return ExperimentalNuageRequest(
idempotency_key="idem-1",
message=ExperimentalNuageMessage(
parts=[ExperimentalNuageTextPart(text="continue from here")]
),
context=ExperimentalNuageContext(
repositories=[
ExperimentalNuageRepository(
repo_url="https://github.com/owner/repo", branch="main"
)
]
),
)
@pytest.mark.asyncio
async def test_start_posts_experimental_nuage_request() -> None:
seen_request: httpx.Request | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_request
seen_request = request
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
experimental_nuage = ExperimentalNuageClient(
"https://chat.example.com/", "api-key", client=client
)
response = await experimental_nuage.start(_request())
assert seen_request is not None
assert str(seen_request.url) == "https://chat.example.com/api/v1/code/sessions"
assert seen_request.headers["authorization"] == "Bearer api-key"
assert seen_request.headers["content-type"] == "application/json"
assert json.loads(seen_request.content) == {
"project_name": "Vibe CLI",
"source": "vibe_code_cli",
"idempotencyKey": "idem-1",
"message": {
"role": "user",
"parts": [{"type": "text", "text": "continue from here"}],
},
"context": {
"repositories": [
{"repoUrl": "https://github.com/owner/repo", "branch": "main"}
]
},
}
assert response.nuage_session_id == "controller-session-id"
assert response.nuage_web_session_id == "web-session-id"
assert response.nuage_project_id == "project-id"
assert response.status == "running"
assert response.url == "https://chat.example.com/code/project-id/web-session-id"
@pytest.mark.asyncio
async def test_start_omits_empty_branch() -> None:
seen_body: dict[str, object] | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
request = ExperimentalNuageRequest(
idempotency_key="idem-1",
message=ExperimentalNuageMessage(
parts=[ExperimentalNuageTextPart(text="prompt")]
),
context=ExperimentalNuageContext(
repositories=[
ExperimentalNuageRepository(repo_url="https://github.com/owner/repo")
]
),
)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
experimental_nuage = ExperimentalNuageClient(
"https://chat.example.com", "api-key", client=client
)
await experimental_nuage.start(request)
assert seen_body == {
"project_name": "Vibe CLI",
"source": "vibe_code_cli",
"idempotencyKey": "idem-1",
"message": {"role": "user", "parts": [{"type": "text", "text": "prompt"}]},
"context": {"repositories": [{"repoUrl": "https://github.com/owner/repo"}]},
}
@pytest.mark.asyncio
async def test_start_serializes_commit_sha_and_diff() -> None:
seen_body: dict[str, object] | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "s",
"webSessionId": "ws",
"projectId": "p",
"status": "running",
"url": "https://chat.example.com/code/p/ws",
},
)
request = ExperimentalNuageRequest(
idempotency_key="idem-1",
message=ExperimentalNuageMessage(
parts=[ExperimentalNuageTextPart(text="prompt")]
),
context=ExperimentalNuageContext(
repositories=[
ExperimentalNuageRepository(
repo_url="https://github.com/owner/repo",
branch="main",
commit_sha="abc123",
diff=ExperimentalNuageDiff(
format="git-diff",
encoding="base64",
compression="zstd",
content="ZGlmZnM=",
),
)
]
),
)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
nuage = ExperimentalNuageClient(
"https://chat.example.com", "api-key", client=client
)
await nuage.start(request)
assert seen_body is not None
repos = seen_body["context"]["repositories"]
assert len(repos) == 1
assert repos[0]["commitSha"] == "abc123"
assert repos[0]["diff"] == {
"format": "git-diff",
"encoding": "base64",
"compression": "zstd",
"content": "ZGlmZnM=",
}
@pytest.mark.asyncio
async def test_start_raises_for_unsuccessful_response() -> None:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(401, text="Unauthorized")
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
experimental_nuage = ExperimentalNuageClient(
"https://chat.example.com", "api-key", client=client
)
with pytest.raises(ServiceTeleportError, match="Nuage start"):
await experimental_nuage.start(_request())
@pytest.mark.asyncio
async def test_start_raises_for_invalid_response() -> None:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"url": "https://chat.example.com/code/1/2"})
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
experimental_nuage = ExperimentalNuageClient(
"https://chat.example.com", "api-key", client=client
)
with pytest.raises(ServiceTeleportError, match="response was invalid"):
await experimental_nuage.start(_request())

View file

@ -1,68 +0,0 @@
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

View file

@ -1,262 +1,189 @@
from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock
import json
import httpx
import pytest
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.teleport.nuage import (
ChatAssistantParams,
GitHubParams,
NuageClient,
WorkflowIntegrations,
WorkflowParams,
NuageContext,
NuageDiff,
NuageMessage,
NuageRepository,
NuageRequest,
NuageTextPart,
)
class TestNuageModels:
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 _request() -> NuageRequest:
return NuageRequest(
idempotency_key="idem-1",
message=NuageMessage(parts=[NuageTextPart(text="continue from here")]),
context=NuageContext(
repositories=[
NuageRepository(repo_url="https://github.com/owner/repo", branch="main")
]
),
)
def test_github_params_with_values(self) -> None:
params = GitHubParams(
repo="owner/repo",
branch="main",
commit="abc123",
pr_number=42,
teleported_diffs=b"base64data",
@pytest.mark.asyncio
async def test_start_posts_nuage_request() -> None:
seen_request: httpx.Request | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_request
seen_request = request
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
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_chat_assistant_params(self) -> None:
params = ChatAssistantParams(
user_message="test message", project_name="test-project"
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
nuage = NuageClient("https://chat.example.com/", "api-key", client=client)
response = await nuage.start(_request())
assert seen_request is not None
assert str(seen_request.url) == "https://chat.example.com/api/v1/code/sessions"
assert seen_request.headers["authorization"] == "Bearer api-key"
assert seen_request.headers["content-type"] == "application/json"
assert json.loads(seen_request.content) == {
"project_name": "Vibe CLI",
"source": "vibe_code_cli",
"idempotencyKey": "idem-1",
"message": {
"role": "user",
"parts": [{"type": "text", "text": "continue from here"}],
},
"context": {
"repositories": [
{"repoUrl": "https://github.com/owner/repo", "branch": "main"}
]
},
}
assert response.nuage_session_id == "controller-session-id"
assert response.nuage_web_session_id == "web-session-id"
assert response.nuage_project_id == "project-id"
assert response.status == "running"
assert response.url == "https://chat.example.com/code/project-id/web-session-id"
@pytest.mark.asyncio
async def test_start_omits_empty_branch() -> None:
seen_body: dict[str, object] | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
assert params.user_message == "test message"
assert params.project_name == "test-project"
def test_workflow_integrations(self) -> None:
integrations = WorkflowIntegrations(
github=GitHubParams(repo="owner/repo"),
chat_assistant=ChatAssistantParams(user_message="test"),
request = NuageRequest(
idempotency_key="idem-1",
message=NuageMessage(parts=[NuageTextPart(text="prompt")]),
context=NuageContext(
repositories=[NuageRepository(repo_url="https://github.com/owner/repo")]
),
)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
nuage = NuageClient("https://chat.example.com", "api-key", client=client)
await nuage.start(request)
assert seen_body == {
"project_name": "Vibe CLI",
"source": "vibe_code_cli",
"idempotencyKey": "idem-1",
"message": {"role": "user", "parts": [{"type": "text", "text": "prompt"}]},
"context": {"repositories": [{"repoUrl": "https://github.com/owner/repo"}]},
}
@pytest.mark.asyncio
async def test_start_serializes_commit_sha_and_diff() -> None:
seen_body: dict[str, object] | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "s",
"webSessionId": "ws",
"projectId": "p",
"status": "running",
"url": "https://chat.example.com/code/p/ws",
},
)
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",
integrations=WorkflowIntegrations(
github=GitHubParams(
repo="owner/repo",
request = NuageRequest(
idempotency_key="idem-1",
message=NuageMessage(parts=[NuageTextPart(text="prompt")]),
context=NuageContext(
repositories=[
NuageRepository(
repo_url="https://github.com/owner/repo",
branch="main",
commit="abc123",
pr_number=42,
teleported_diffs=b"base64data",
),
chat_assistant=ChatAssistantParams(user_message="test"),
),
)
data = params.model_dump()
assert data["prompt"] == "test prompt"
assert data["integrations"]["github"]["repo"] == "owner/repo"
assert data["integrations"]["github"]["pr_number"] == 42
assert data["integrations"]["github"]["teleported_diffs"] == b"base64data"
commit_sha="abc123",
diff=NuageDiff(
format="git-diff",
encoding="base64",
compression="zstd",
content="ZGlmZnM=",
),
)
]
),
)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
nuage = NuageClient("https://chat.example.com", "api-key", client=client)
await nuage.start(request)
assert seen_body is not None
repos = seen_body["context"]["repositories"]
assert len(repos) == 1
assert repos[0]["commitSha"] == "abc123"
assert repos[0]["diff"] == {
"format": "git-diff",
"encoding": "base64",
"compression": "zstd",
"content": "ZGlmZnM=",
}
class TestNuageClientContextManager:
@pytest.mark.asyncio
async def test_creates_client_on_enter(self) -> None:
nuage = NuageClient("https://api.example.com", "api-key", "workflow-id")
assert nuage._client is None
async with nuage:
assert nuage._client is not None
assert nuage._client is None
@pytest.mark.asyncio
async def test_start_raises_for_unsuccessful_response() -> None:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(401, text="Unauthorized")
@pytest.mark.asyncio
async def test_uses_provided_client(self) -> None:
external_client = httpx.AsyncClient()
nuage = NuageClient(
"https://api.example.com", "api-key", "workflow-id", client=external_client
)
async with nuage:
assert nuage._client is external_client
assert nuage._client is external_client
await external_client.aclose()
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
nuage = NuageClient("https://chat.example.com", "api-key", client=client)
with pytest.raises(ServiceTeleportError, match="Nuage start"):
await nuage.start(_request())
class TestNuageClientStartWorkflow:
@pytest.fixture
def mock_client(self) -> MagicMock:
client = MagicMock(spec=httpx.AsyncClient)
return client
@pytest.mark.asyncio
async def test_start_raises_for_invalid_response() -> None:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"url": "https://chat.example.com/code/1/2"})
@pytest.fixture
def nuage(self, mock_client: MagicMock) -> NuageClient:
return NuageClient(
"https://api.example.com", "api-key", "workflow-id", client=mock_client
)
@pytest.mark.asyncio
async def test_start_workflow_success(
self, nuage: NuageClient, mock_client: MagicMock
) -> None:
mock_response = MagicMock()
mock_response.is_success = True
mock_response.json.return_value = {"execution_id": "exec-123"}
mock_client.post = AsyncMock(return_value=mock_response)
params = WorkflowParams(prompt="test", integrations=WorkflowIntegrations())
execution_id = await nuage.start_workflow(params)
assert execution_id == "exec-123"
mock_client.post.assert_called_once()
call_args = mock_client.post.call_args
assert "workflows/workflow-id/execute" in call_args[0][0]
@pytest.mark.asyncio
async def test_start_workflow_failure(
self, nuage: NuageClient, mock_client: MagicMock
) -> None:
mock_response = MagicMock()
mock_response.is_success = False
mock_response.text = "Internal Server Error"
mock_client.post = AsyncMock(return_value=mock_response)
params = WorkflowParams(prompt="test", integrations=WorkflowIntegrations())
with pytest.raises(
ServiceTeleportError, match="Vibe Code workflow trigger failed"
):
await nuage.start_workflow(params)
class TestNuageClientGetGitHubIntegration:
@pytest.fixture
def mock_client(self) -> MagicMock:
return MagicMock(spec=httpx.AsyncClient)
@pytest.fixture
def nuage(self, mock_client: MagicMock) -> NuageClient:
return NuageClient(
"https://api.example.com", "api-key", "workflow-id", client=mock_client
)
@pytest.mark.asyncio
async def test_get_github_integration_connected(
self, nuage: NuageClient, mock_client: MagicMock
) -> None:
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)
result = await nuage.get_github_integration("exec-123")
assert result.connected is True
assert result.oauth_url is None
@pytest.mark.asyncio
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()
mock_response.is_success = False
mock_response.text = "Not found"
mock_client.post = AsyncMock(return_value=mock_response)
with pytest.raises(
ServiceTeleportError, match="Failed to get GitHub integration"
):
await nuage.get_github_integration("exec-123")
class TestNuageClientGetChatAssistantUrl:
@pytest.fixture
def mock_client(self) -> MagicMock:
return MagicMock(spec=httpx.AsyncClient)
@pytest.fixture
def nuage(self, mock_client: MagicMock) -> NuageClient:
return NuageClient(
"https://api.example.com", "api-key", "workflow-id", client=mock_client
)
@pytest.mark.asyncio
async def test_get_chat_assistant_url_success(
self, nuage: NuageClient, mock_client: MagicMock
) -> None:
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)
url = await nuage.get_chat_assistant_url("exec-123")
assert url == "https://chat.example.com/thread/123"
@pytest.mark.asyncio
async def test_get_chat_assistant_url_none(
self, nuage: NuageClient, mock_client: MagicMock
) -> None:
mock_response = MagicMock()
mock_response.is_success = True
mock_response.json.return_value = {"result": {"chat_url": None}}
mock_client.post = AsyncMock(return_value=mock_response)
url = await nuage.get_chat_assistant_url("exec-123")
assert url is None
@pytest.mark.asyncio
async def test_get_chat_assistant_url_failure(
self, nuage: NuageClient, mock_client: MagicMock
) -> None:
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 get chat assistant integration"
):
await nuage.get_chat_assistant_url("exec-123")
class TestNuageClientHeaders:
def test_headers_include_auth(self) -> None:
nuage = NuageClient("https://api.example.com", "test-api-key", "workflow-id")
headers = nuage._headers()
assert headers["Authorization"] == "Bearer test-api-key"
assert headers["Content-Type"] == "application/json"
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
nuage = NuageClient("https://chat.example.com", "api-key", client=client)
with pytest.raises(ServiceTeleportError, match="response was invalid"):
await nuage.start(_request())

View file

@ -13,25 +13,19 @@ import httpx
import pytest
import zstandard
from vibe.core.config import VibeConfig
from vibe.core.teleport.errors import (
ServiceTeleportError,
ServiceTeleportNotSupportedError,
)
from vibe.core.teleport.git import GitRepoInfo
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,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
)
@ -42,17 +36,38 @@ def _reimport_agent_loop() -> Any:
return importlib.import_module("vibe.core.agent_loop")
def _make_service(tmp_path: Path, **kwargs: Any) -> TeleportService:
return TeleportService(
session_logger=MagicMock(),
vibe_code_sessions_base_url=kwargs.pop(
"vibe_code_sessions_base_url", "https://api.example.com"
),
vibe_code_api_key=kwargs.pop("vibe_code_api_key", "api-key"),
workdir=tmp_path,
**kwargs,
)
def _mock_handler() -> Any:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
return handler
class TestTeleportServiceCompressDiff:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
return _make_service(tmp_path)
def test_returns_none_for_empty_diff(self, service: TeleportService) -> None:
assert service._compress_diff("") is None
@ -72,84 +87,20 @@ class TestTeleportServiceCompressDiff:
service._compress_diff(large_diff, max_size=100)
class TestTeleportServiceBuildGitHubParams:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
def test_builds_params_from_git_info(self, service: TeleportService) -> None:
git_info = GitRepoInfo(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
params = service._build_github_params(git_info)
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(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="diff content",
)
params = service._build_github_params(git_info)
assert params.teleported_diffs is not None
class TestTeleportServiceValidateConfig:
def test_raises_when_no_api_key(self, tmp_path: Path) -> None:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="",
workdir=tmp_path,
)
service = _make_service(tmp_path, vibe_code_api_key="")
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:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="valid-key",
workdir=tmp_path,
)
service = _make_service(tmp_path, vibe_code_api_key="valid-key")
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.vibe_code_api_key_env_var = "CUSTOM_API_KEY"
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="",
workdir=tmp_path,
vibe_config=mock_config,
)
service = _make_service(tmp_path, vibe_code_api_key="", vibe_config=mock_config)
with pytest.raises(ServiceTeleportError, match="CUSTOM_API_KEY not set"):
service._validate_config()
@ -157,14 +108,7 @@ class TestTeleportServiceValidateConfig:
class TestTeleportServiceCheckSupported:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
return _make_service(tmp_path)
@pytest.mark.asyncio
async def test_check_supported_calls_git_info(
@ -189,7 +133,7 @@ class TestTeleportServiceCheckSupported:
) -> None:
service._git.get_info = AsyncMock(
side_effect=ServiceTeleportNotSupportedError(
"Teleport requires a git repository. cd into a project with a .git directory and try again."
"Teleport requires a git repository."
)
)
with pytest.raises(ServiceTeleportNotSupportedError):
@ -199,14 +143,7 @@ class TestTeleportServiceCheckSupported:
class TestTeleportServiceIsSupported:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
return _make_service(tmp_path)
@pytest.mark.asyncio
async def test_is_supported_returns_true(self, service: TeleportService) -> None:
@ -220,248 +157,8 @@ class TestTeleportServiceIsSupported:
class TestTeleportServiceExecute:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
service._git.fetch = AsyncMock()
service._git.is_branch_pushed = AsyncMock(return_value=True)
return service
@pytest.fixture
def git_info(self) -> GitRepoInfo:
return GitRepoInfo(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
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_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)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
events = []
gen = service.execute("test prompt", session)
async for event in gen:
events.append(event)
assert isinstance(events[0], TeleportCheckingGitEvent)
assert isinstance(events[1], TeleportStartingWorkflowEvent)
assert isinstance(events[2], TeleportWaitingForGitHubEvent)
assert isinstance(events[3], TeleportAuthCompleteEvent)
assert isinstance(events[4], TeleportFetchingUrlEvent)
assert isinstance(events[5], TeleportCompleteEvent)
assert events[5].url == "https://chat.example.com/123"
workflow_params = mock_nuage.start_workflow.call_args.args[0]
assert workflow_params.integrations.chat_assistant is not None
assert workflow_params.integrations.chat_assistant.project_name is None
@pytest.mark.asyncio
async def test_execute_requires_push_and_user_approves(
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)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
events = []
gen = service.execute("test prompt", session)
event = await gen.asend(None)
events.append(event)
assert isinstance(event, TeleportCheckingGitEvent)
event = await gen.asend(None)
events.append(event)
assert isinstance(event, TeleportPushRequiredEvent)
assert event.unpushed_count == 3
event = await gen.asend(TeleportPushResponseEvent(approved=True))
events.append(event)
assert isinstance(event, TeleportPushingEvent)
async for event in gen:
events.append(event)
assert isinstance(events[-1], TeleportCompleteEvent)
@pytest.mark.asyncio
async def test_execute_requires_push_and_user_declines(
self, service: TeleportService, git_info: GitRepoInfo
) -> 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=1)
session = TeleportSession()
gen = service.execute("test prompt", session)
await gen.asend(None)
await gen.asend(None)
with pytest.raises(ServiceTeleportError, match="Teleport cancelled"):
await gen.asend(TeleportPushResponseEvent(approved=False))
@pytest.mark.asyncio
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)
github_pending = MagicMock()
github_pending.connected = False
github_pending.oauth_url = "https://github.com/login/oauth"
github_pending.error = "Please connect GitHub"
github_pending.status = GitHubStatus.WAITING_FOR_OAUTH
github_connected = MagicMock()
github_connected.connected = True
github_connected.oauth_url = None
github_connected.error = None
github_connected.status = GitHubStatus.CONNECTED
async def _oauth_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield github_pending
yield github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _oauth_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
events = []
gen = service.execute("test prompt", session)
async for event in gen:
events.append(event)
assert isinstance(events[0], TeleportCheckingGitEvent)
assert isinstance(events[1], TeleportStartingWorkflowEvent)
assert isinstance(events[2], TeleportWaitingForGitHubEvent)
assert events[2].message is None
assert isinstance(events[3], TeleportAuthRequiredEvent)
assert events[3].oauth_url == "https://github.com/login/oauth"
assert isinstance(events[4], TeleportWaitingForGitHubEvent)
assert events[4].message == "Please connect GitHub"
assert isinstance(events[5], TeleportAuthCompleteEvent)
assert isinstance(events[-1], TeleportCompleteEvent)
@pytest.mark.asyncio
async def test_execute_raises_when_chat_url_is_none(
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)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(return_value=None)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
gen = service.execute("test prompt", session)
with pytest.raises(ServiceTeleportError, match="not available"):
async for _ in gen:
pass
@pytest.mark.asyncio
async def test_execute_uses_default_prompt_when_none(
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)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
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 "teleported" in call_args[0][0].prompt.lower()
@pytest.mark.asyncio
async def test_execute_uses_experimental_nuage_when_enabled(
self, tmp_path: Path
) -> None:
async def test_execute_happy_path(self, tmp_path: Path) -> None:
seen_body: dict[str, object] | None = None
seen_url: str | None = None
@ -480,15 +177,10 @@ class TestTeleportServiceExecute:
},
)
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://chat.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
service = _make_service(
tmp_path,
vibe_code_sessions_base_url="https://chat.example.com",
client=client,
)
service._git.fetch = AsyncMock()
@ -505,10 +197,7 @@ class TestTeleportServiceExecute:
service._git.is_commit_pushed = AsyncMock(return_value=True)
service._git.is_branch_pushed = AsyncMock(return_value=True)
events = [
event
async for event in service.execute("test prompt", TeleportSession())
]
events = [event async for event in service.execute("test prompt")]
assert isinstance(events[0], TeleportCheckingGitEvent)
assert isinstance(events[1], TeleportStartingWorkflowEvent)
@ -516,7 +205,6 @@ class TestTeleportServiceExecute:
assert (
events[2].url == "https://chat.example.com/code/project-id/web-session-id"
)
assert service._nuage_client_instance is None
assert seen_url == "https://chat.example.com/api/v1/code/sessions"
assert seen_body is not None
assert seen_body["message"] == {
@ -535,75 +223,8 @@ class TestTeleportServiceExecute:
assert "idempotencyKey" in seen_body
@pytest.mark.asyncio
async def test_execute_experimental_nuage_uses_last_message_when_prompt_missing(
self, tmp_path: Path
) -> None:
seen_body: dict[str, object] | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
client=client,
)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
)
service._git.is_commit_pushed = AsyncMock(return_value=True)
service._git.is_branch_pushed = AsyncMock(return_value=True)
session = TeleportSession(
messages=[{"role": "user", "content": "help me refactor"}]
)
events = [event async for event in service.execute(None, session)]
assert isinstance(events[-1], TeleportCompleteEvent)
assert seen_body is not None
assert seen_body["message"] == {
"role": "user",
"parts": [{"type": "text", "text": "help me refactor (continue)"}],
}
@pytest.mark.asyncio
async def test_execute_experimental_nuage_requires_branch(
self, tmp_path: Path
) -> None:
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
)
async def test_execute_requires_branch(self, tmp_path: Path) -> None:
service = _make_service(tmp_path)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
@ -617,38 +238,24 @@ class TestTeleportServiceExecute:
)
with pytest.raises(ServiceTeleportError, match="checked-out branch"):
async for _ in service.execute("test prompt", TeleportSession()):
async for _ in service.execute("test prompt"):
pass
service._git.fetch.assert_not_awaited()
@pytest.mark.asyncio
async def test_execute_experimental_nuage_keeps_push_confirmation(
self, tmp_path: Path
) -> None:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
async def test_execute_rejects_empty_prompt(self, tmp_path: Path) -> None:
service = _make_service(tmp_path)
with pytest.raises(ServiceTeleportError, match="non-empty prompt"):
async for _ in service.execute(""):
pass
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
client=client,
)
@pytest.mark.asyncio
async def test_execute_push_confirmation_approved(self, tmp_path: Path) -> None:
async with httpx.AsyncClient(
transport=httpx.MockTransport(_mock_handler())
) as client:
service = _make_service(tmp_path, client=client)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
@ -665,7 +272,7 @@ class TestTeleportServiceExecute:
service._git.get_unpushed_commit_count = AsyncMock(return_value=2)
service._git.push_current_branch = AsyncMock(return_value=True)
gen = service.execute("test prompt", TeleportSession())
gen = service.execute("test prompt")
assert isinstance(await gen.asend(None), TeleportCheckingGitEvent)
push_event = await gen.asend(None)
assert isinstance(push_event, TeleportPushRequiredEvent)
@ -681,18 +288,36 @@ class TestTeleportServiceExecute:
assert isinstance(events[0], TeleportStartingWorkflowEvent)
assert isinstance(events[1], TeleportCompleteEvent)
@pytest.mark.asyncio
async def test_execute_push_confirmation_declined(self, tmp_path: Path) -> None:
service = _make_service(tmp_path)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
)
service._git.is_commit_pushed = AsyncMock(return_value=False)
service._git.is_branch_pushed = AsyncMock(return_value=True)
service._git.get_unpushed_commit_count = AsyncMock(return_value=1)
gen = service.execute("test prompt")
await gen.asend(None)
await gen.asend(None)
with pytest.raises(ServiceTeleportError, match="Teleport cancelled"):
await gen.asend(TeleportPushResponseEvent(approved=False))
class TestTeleportServiceContextManager:
@pytest.mark.asyncio
async def test_creates_client_on_enter(self, tmp_path: Path) -> None:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
service = _make_service(tmp_path)
assert service._client is None
async with service:
assert service._client is not None

View file

@ -10,15 +10,12 @@ from vibe.core.agent_loop import AgentLoop, TeleportError
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.teleport.teleport import TeleportService
from vibe.core.teleport.types import (
TeleportAuthRequiredEvent,
TeleportCheckingGitEvent,
TeleportCompleteEvent,
TeleportFetchingUrlEvent,
TeleportPushingEvent,
TeleportPushRequiredEvent,
TeleportPushResponseEvent,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
)
from vibe.core.types import LLMMessage, Role
@ -46,9 +43,6 @@ class TestTeleportAgentLoopTelemetry:
yield TeleportPushRequiredEvent()
yield TeleportPushingEvent()
yield TeleportStartingWorkflowEvent()
yield TeleportWaitingForGitHubEvent()
yield TeleportAuthRequiredEvent(oauth_url="https://github.com/auth")
yield TeleportFetchingUrlEvent()
yield TeleportCompleteEvent(url="https://chat.example.com/123")
agent_loop.messages.append(LLMMessage(role=Role.user, content="hello"))
@ -73,7 +67,6 @@ class TestTeleportAgentLoopTelemetry:
assert telemetry_events[-1]["event_name"] == "vibe.teleport_completed"
assert telemetry_events[-1]["properties"] == {
"push_required": True,
"github_auth_required": True,
"nb_session_messages": 1,
"session_id": agent_loop.session_id,
}
@ -109,7 +102,6 @@ class TestTeleportAgentLoopTelemetry:
"stage": "workflow_start",
"error_class": "ServiceTeleportError",
"push_required": False,
"github_auth_required": False,
"nb_session_messages": 1,
"session_id": agent_loop.session_id,
}
@ -154,7 +146,6 @@ class TestTeleportAgentLoopTelemetry:
"stage": "cancelled",
"error_class": "ServiceTeleportError",
"push_required": True,
"github_auth_required": False,
"nb_session_messages": 1,
"session_id": agent_loop.session_id,
}
@ -190,7 +181,6 @@ class TestTeleportAgentLoopTelemetry:
"stage": "cancelled",
"error_class": "CancelledError",
"push_required": False,
"github_auth_required": False,
"nb_session_messages": 1,
"session_id": agent_loop.session_id,
}
@ -225,7 +215,6 @@ class TestTeleportAgentLoopTelemetry:
"stage": "cancelled",
"error_class": "CancelledError",
"push_required": False,
"github_auth_required": False,
"nb_session_messages": 1,
"session_id": agent_loop.session_id,
}

View file

@ -14,7 +14,6 @@ import time
from typing import Any, ClassVar, assert_never, cast
from uuid import uuid4
from weakref import WeakKeyDictionary
import webbrowser
from pydantic import BaseModel
from rich import print as rprint
@ -160,16 +159,12 @@ from vibe.core.session.title_format import format_session_title
from vibe.core.skills.manager import SkillManager
from vibe.core.teleport.telemetry import send_teleport_early_failure_telemetry
from vibe.core.teleport.types import (
TeleportAuthCompleteEvent,
TeleportAuthRequiredEvent,
TeleportCheckingGitEvent,
TeleportCompleteEvent,
TeleportFetchingUrlEvent,
TeleportPushingEvent,
TeleportPushRequiredEvent,
TeleportPushResponseEvent,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
)
from vibe.core.tools.builtins.ask_user_question import (
AskUserQuestionArgs,
@ -1604,15 +1599,6 @@ class VibeApp(App): # noqa: PLR0904
teleport_msg.set_status("Syncing with remote...")
case TeleportStartingWorkflowEvent():
teleport_msg.set_status("Teleporting...")
case TeleportWaitingForGitHubEvent(message=msg):
teleport_msg.set_status(msg or "Connecting to GitHub...")
case TeleportAuthRequiredEvent(oauth_url=url, message=msg):
webbrowser.open(url)
teleport_msg.set_status(msg or "Authorizing GitHub...")
case TeleportAuthCompleteEvent():
teleport_msg.set_status("GitHub authorized")
case TeleportFetchingUrlEvent():
teleport_msg.set_status("Finalizing...")
case TeleportCompleteEvent(url=url):
teleport_msg.set_complete(url)
except TeleportError as e:

View file

@ -675,10 +675,8 @@ class AgentLoop: # noqa: PLR0904
raise TeleportError("_TeleportService is unexpectedly None")
self._teleport_service = _TeleportService(
session_logger=self.session_logger,
vibe_code_base_url=self.config.vibe_code_base_url,
vibe_code_workflow_id=self.config.vibe_code_workflow_id,
vibe_code_sessions_base_url=self.config.vibe_code_sessions_base_url,
vibe_code_api_key=self.config.vibe_code_api_key,
vibe_code_task_queue=self.config.vibe_code_task_queue,
vibe_config=self._base_config,
)
return self._teleport_service
@ -687,29 +685,23 @@ class AgentLoop: # noqa: PLR0904
async def teleport_to_vibe_code(
self, prompt: str | None
) -> AsyncGenerator[TeleportYieldEvent, TeleportPushResponseEvent | None]:
from vibe.core.teleport.nuage import TeleportSession
session_messages = [
msg.model_dump(exclude_none=True) for msg in self.messages[1:]
]
nb_session_messages = max(len(self.messages) - 1, 0)
if prompt:
resolved_prompt = prompt
else:
last = self._last_user_message()
content = last.content if last else None
resolved_prompt = (
f"{content} (continue)" if isinstance(content, str) and content else ""
)
telemetry_tracker = TeleportTelemetryTracker(
telemetry_client=self.telemetry_client,
nb_session_messages=len(session_messages),
stage="no_history"
if prompt is None and not session_messages
else "git_check",
)
session = TeleportSession(
metadata={
"agent": self.agent_profile.name,
"model": self.config.active_model,
"stats": self.stats.model_dump(),
},
messages=session_messages,
nb_session_messages=nb_session_messages,
stage="no_history" if not resolved_prompt else "git_check",
)
try:
async with self.teleport_service:
gen = self.teleport_service.execute(prompt=prompt, session=session)
gen = self.teleport_service.execute(prompt=resolved_prompt)
response: TeleportPushResponseEvent | None = None
while True:
try:
@ -733,6 +725,16 @@ class AgentLoop: # noqa: PLR0904
telemetry_tracker.send_failure_if_needed()
self._teleport_service = None
def _last_user_message(self) -> LLMMessage | None:
return next(
(
m
for m in reversed(self.messages)
if m.role == Role.user and not m.injected
),
None,
)
def _setup_middleware(self) -> None:
"""Configure middleware pipeline for this conversation."""
self.middleware_pipeline.clear()
@ -1330,14 +1332,7 @@ class AgentLoop: # noqa: PLR0904
available_tools = self.format_handler.get_available_tools(self.tool_manager)
tool_choice = self.format_handler.get_tool_choice()
last_user_message = next(
(
m
for m in reversed(self.messages)
if m.role == Role.user and not m.injected
),
None,
)
last_user_message = self._last_user_message()
self.telemetry_client.send_request_sent(
model=active_model.alias,
nb_context_chars=sum(len(m.content or "") for m in self.messages),
@ -1400,14 +1395,7 @@ class AgentLoop: # noqa: PLR0904
available_tools = self.format_handler.get_available_tools(self.tool_manager)
tool_choice = self.format_handler.get_tool_choice()
last_user_message = next(
(
m
for m in reversed(self.messages)
if m.role == Role.user and not m.injected
),
None,
)
last_user_message = self._last_user_message()
self.telemetry_client.send_request_sent(
model=active_model.alias,
nb_context_chars=sum(len(m.content or "") for m in self.messages),

View file

@ -525,11 +525,12 @@ class VibeConfig(BaseSettings):
vibe_code_enabled: bool = Field(default=True, exclude=True)
vibe_code_base_url: str = Field(default="https://api.mistral.ai", exclude=True)
vibe_code_sessions_base_url: str = Field(
default="https://chat.mistral.ai", exclude=True
)
vibe_code_workflow_id: str = Field(default="__shared-nuage-workflow", exclude=True)
vibe_code_task_queue: str | None = Field(default="shared-vibe-nuage", exclude=True)
vibe_code_api_key_env_var: str = Field(default="MISTRAL_API_KEY", exclude=True)
vibe_code_project_name: str | None = Field(default=None, exclude=True)
vibe_code_experimental_nuage_enabled: bool = Field(default=False, exclude=True)
# TODO(otel): remove exclude=True once the feature is publicly available
enable_otel: bool = Field(default=False, exclude=True)

View file

@ -6,15 +6,11 @@ import sys
from typing import TextIO
from vibe.core.teleport.types import (
TeleportAuthCompleteEvent,
TeleportAuthRequiredEvent,
TeleportCheckingGitEvent,
TeleportCompleteEvent,
TeleportFetchingUrlEvent,
TeleportPushingEvent,
TeleportPushRequiredEvent,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
)
from vibe.core.types import AssistantEvent, BaseEvent, LLMMessage, OutputFormat
@ -62,14 +58,6 @@ class TextOutputFormatter(OutputFormatter):
self._print("Syncing with remote...")
case TeleportStartingWorkflowEvent():
self._print("Teleporting...")
case TeleportWaitingForGitHubEvent(message=msg):
self._print(msg or "Connecting to GitHub...")
case TeleportAuthRequiredEvent(oauth_url=url, message=msg):
self._print(msg or f"Open to authorize GitHub: {url}")
case TeleportAuthCompleteEvent():
self._print("GitHub authorized")
case TeleportFetchingUrlEvent():
self._print("Finalizing...")
case TeleportCompleteEvent():
self._final_response = event.url

View file

@ -349,15 +349,10 @@ class TelemetryClient:
)
def send_teleport_completed(
self,
*,
push_required: bool,
github_auth_required: bool,
nb_session_messages: int,
self, *, push_required: bool, nb_session_messages: int
) -> None:
payload: TeleportCompletedPayload = {
"push_required": push_required,
"github_auth_required": github_auth_required,
"nb_session_messages": nb_session_messages,
}
self.send_telemetry_event("vibe.teleport_completed", dict(payload))
@ -368,14 +363,12 @@ class TelemetryClient:
stage: TeleportFailureStage,
error_class: str,
push_required: bool,
github_auth_required: bool,
nb_session_messages: int,
) -> None:
payload: TeleportFailedPayload = {
"stage": stage,
"error_class": error_class,
"push_required": push_required,
"github_auth_required": github_auth_required,
"nb_session_messages": nb_session_messages,
}
self.send_telemetry_event("vibe.teleport_failed", dict(payload))

View file

@ -40,20 +40,12 @@ class TelemetryRequestMetadata(TelemetryBaseMetadata):
TeleportFailureStage = Literal[
"no_history",
"remote_session",
"git_check",
"push",
"workflow_start",
"github_auth",
"fetch_url",
"cancelled",
"no_history", "remote_session", "git_check", "push", "workflow_start", "cancelled"
]
class TeleportCompletedPayload(TypedDict):
push_required: bool
github_auth_required: bool
nb_session_messages: int

View file

@ -1,132 +0,0 @@
from __future__ import annotations
import types
from typing import Literal
import httpx
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.utils.http import build_ssl_context
class ExperimentalNuageTextPart(BaseModel):
model_config = ConfigDict(extra="forbid")
type: str = "text"
text: str
class ExperimentalNuageMessage(BaseModel):
model_config = ConfigDict(extra="forbid")
role: str = "user"
parts: list[ExperimentalNuageTextPart]
class ExperimentalNuageDiff(BaseModel):
model_config = ConfigDict(extra="forbid")
format: Literal["git-diff"] = "git-diff"
encoding: Literal["base64"] = "base64"
compression: Literal["zstd"] = "zstd"
content: str
class ExperimentalNuageRepository(BaseModel):
model_config = ConfigDict(extra="forbid")
repo_url: str = Field(serialization_alias="repoUrl")
branch: str | None = None
commit_sha: str | None = Field(default=None, serialization_alias="commitSha")
diff: ExperimentalNuageDiff | None = None
class ExperimentalNuageContext(BaseModel):
model_config = ConfigDict(extra="forbid")
repositories: list[ExperimentalNuageRepository]
class ExperimentalNuageRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
project_name: str = Field(default="Vibe CLI", serialization_alias="project_name")
source: str = "vibe_code_cli"
idempotency_key: str = Field(serialization_alias="idempotencyKey")
message: ExperimentalNuageMessage
context: ExperimentalNuageContext
class ExperimentalNuageResponse(BaseModel):
model_config = ConfigDict(extra="ignore")
nuage_session_id: str = Field(validation_alias="sessionId")
nuage_web_session_id: str = Field(validation_alias="webSessionId")
nuage_project_id: str = Field(validation_alias="projectId")
status: str
url: str
class ExperimentalNuageClient:
def __init__(
self,
base_url: str,
api_key: str,
*,
client: httpx.AsyncClient | None = None,
timeout: float = 60.0,
) -> None:
self._base_url = base_url.rstrip("/")
self._api_key = api_key
self._client = client
self._owns_client = client is None
self._timeout = timeout
async def __aenter__(self) -> ExperimentalNuageClient:
if self._client is None:
self._client = httpx.AsyncClient(
timeout=httpx.Timeout(self._timeout), verify=build_ssl_context()
)
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: types.TracebackType | None,
) -> None:
if self._owns_client and self._client:
await self._client.aclose()
self._client = None
@property
def _http_client(self) -> httpx.AsyncClient:
if self._client is None:
self._client = httpx.AsyncClient(
timeout=httpx.Timeout(self._timeout), verify=build_ssl_context()
)
self._owns_client = True
return self._client
def _headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self._api_key}",
"Content-Type": "application/json",
}
async def start(
self, request: ExperimentalNuageRequest
) -> ExperimentalNuageResponse:
response = await self._http_client.post(
f"{self._base_url}/api/v1/code/sessions",
headers=self._headers(),
json=request.model_dump(mode="json", by_alias=True, exclude_none=True),
)
if not response.is_success:
raise ServiceTeleportError(f"Vibe Code Nuage start failed: {response.text}")
try:
return ExperimentalNuageResponse.model_validate(response.json())
except (ValueError, ValidationError) as e:
raise ServiceTeleportError("Vibe Code Nuage response was invalid") from e

View file

@ -1,104 +1,71 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator
from enum import StrEnum, auto
import time
import types
from typing import Any
from typing import Literal
import httpx
from pydantic import BaseModel, Field, ValidationError
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.utils.http import build_ssl_context
class GitHubParams(BaseModel):
repo: str | None = None
branch: str | None = None
commit: str | None = None
pr_number: int | None = None
teleported_diffs: bytes | None = None
class NuageTextPart(BaseModel):
model_config = ConfigDict(extra="forbid")
class ChatAssistantParams(BaseModel):
create_thread: bool = False
user_message: str | None = None
project_name: str | None = None
class TeleportSession(BaseModel):
metadata: dict[str, Any] = Field(default_factory=dict)
messages: list[dict[str, Any]] = Field(default_factory=list)
class WorkflowIntegrations(BaseModel):
github: GitHubParams | None = None
chat_assistant: ChatAssistantParams | None = None
class VibeAgent(BaseModel):
polymorphic_type: str = "vibe_agent"
name: str = "vibe-agent"
vibe_config: dict[str, Any] | None = None
session: TeleportSession | None = None
class WorkflowConfig(BaseModel):
agent: VibeAgent = Field(default_factory=VibeAgent)
class TextChunk(BaseModel):
type: str = "text"
text: str
class WorkflowParams(BaseModel):
prompt: str
message: list[TextChunk] | None = None
config: WorkflowConfig = Field(default_factory=WorkflowConfig)
integrations: WorkflowIntegrations = Field(default_factory=WorkflowIntegrations)
class NuageMessage(BaseModel):
model_config = ConfigDict(extra="forbid")
role: str = "user"
parts: list[NuageTextPart]
class WorkflowExecuteResponse(BaseModel):
execution_id: str
class NuageDiff(BaseModel):
model_config = ConfigDict(extra="forbid")
format: Literal["git-diff"] = "git-diff"
encoding: Literal["base64"] = "base64"
compression: Literal["zstd"] = "zstd"
content: str
class GitHubStatus(StrEnum):
PENDING = auto()
WAITING_FOR_OAUTH = auto()
CONNECTED = auto()
OAUTH_TIMEOUT = auto()
ERROR = auto()
class NuageRepository(BaseModel):
model_config = ConfigDict(extra="forbid")
repo_url: str = Field(serialization_alias="repoUrl")
branch: str | None = None
commit_sha: str | None = Field(default=None, serialization_alias="commitSha")
diff: NuageDiff | None = None
class GitHubPublicData(BaseModel):
status: GitHubStatus
oauth_url: str | None = None
error: str | None = None
working_branch: str | None = None
repo: str | None = None
class NuageContext(BaseModel):
model_config = ConfigDict(extra="forbid")
@property
def connected(self) -> bool:
return self.status == GitHubStatus.CONNECTED
@property
def is_error(self) -> bool:
return self.status in {GitHubStatus.OAUTH_TIMEOUT, GitHubStatus.ERROR}
repositories: list[NuageRepository]
class ChatAssistantPublicData(BaseModel):
chat_url: str | None = None
class NuageRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
project_name: str = Field(default="Vibe CLI", serialization_alias="project_name")
source: str = "vibe_code_cli"
idempotency_key: str = Field(serialization_alias="idempotencyKey")
message: NuageMessage
context: NuageContext
class GetChatAssistantIntegrationResponse(BaseModel):
result: ChatAssistantPublicData
class NuageResponse(BaseModel):
model_config = ConfigDict(extra="ignore")
class GetGitHubIntegrationResponse(BaseModel):
result: GitHubPublicData
nuage_session_id: str = Field(validation_alias="sessionId")
nuage_web_session_id: str = Field(validation_alias="webSessionId")
nuage_project_id: str = Field(validation_alias="projectId")
status: str
url: str
class NuageClient:
@ -106,16 +73,12 @@ class NuageClient:
self,
base_url: str,
api_key: str,
workflow_id: str,
*,
task_queue: str | None = None,
client: httpx.AsyncClient | None = None,
timeout: float = 60.0,
) -> None:
self._base_url = base_url.rstrip("/")
self._api_key = api_key
self._workflow_id = workflow_id
self._task_queue = task_queue
self._client = client
self._owns_client = client is None
self._timeout = timeout
@ -152,74 +115,16 @@ class NuageClient:
"Content-Type": "application/json",
}
async def start_workflow(self, params: WorkflowParams) -> str:
async def start(self, request: NuageRequest) -> NuageResponse:
response = await self._http_client.post(
f"{self._base_url}/v1/workflows/{self._workflow_id}/execute",
f"{self._base_url}/api/v1/code/sessions",
headers=self._headers(),
json={
"input": params.model_dump(mode="json"),
"task_queue": self._task_queue,
},
json=request.model_dump(mode="json", by_alias=True, exclude_none=True),
)
if not response.is_success:
error_msg = f"Vibe Code workflow trigger failed: {response.text}"
raise ServiceTeleportError(error_msg)
result = WorkflowExecuteResponse.model_validate(response.json())
return result.execution_id
raise ServiceTeleportError(f"Vibe Code Nuage start failed: {response.text}")
async def get_github_integration(self, execution_id: str) -> GitHubPublicData:
response = await self._http_client.post(
f"{self._base_url}/v1/workflows/executions/{execution_id}/updates",
headers=self._headers(),
json={"name": "get_integration", "input": {"integration_id": "github"}},
)
if not response.is_success:
raise ServiceTeleportError(
f"Failed to get GitHub integration: {response.text}"
)
try:
result = GetGitHubIntegrationResponse.model_validate(response.json())
except ValidationError as e:
data = response.json()
error = data.get("result", {}).get("error")
status = data.get("result", {}).get("status")
raise ServiceTeleportError(
f"GitHub integration error: {error or status}"
) from e
return result.result
async def wait_for_github_connection(
self, execution_id: str, timeout: float = 600.0, interval: float = 2.0
) -> AsyncGenerator[GitHubPublicData, None]:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
github_data = await self.get_github_integration(execution_id)
yield github_data
if github_data.connected:
return
if github_data.is_error:
raise ServiceTeleportError(
github_data.error
or f"GitHub integration failed: {github_data.status.value}"
)
remaining = deadline - time.monotonic()
if remaining <= 0:
break
await asyncio.sleep(min(interval, remaining))
raise ServiceTeleportError("GitHub connection timed out")
async def get_chat_assistant_url(self, execution_id: str) -> str | None:
response = await self._http_client.post(
f"{self._base_url}/v1/workflows/executions/{execution_id}/updates",
headers=self._headers(),
json={
"name": "get_integration",
"input": {"integration_id": "chat_assistant"},
},
)
if not response.is_success:
raise ServiceTeleportError(
f"Failed to get chat assistant integration: {response.text}"
)
result = GetChatAssistantIntegrationResponse.model_validate(response.json())
return result.result.chat_url
return NuageResponse.model_validate(response.json())
except (ValueError, ValidationError) as e:
raise ServiceTeleportError("Vibe Code Nuage response was invalid") from e

View file

@ -6,14 +6,11 @@ from vibe.core.telemetry.send import TelemetryClient
from vibe.core.telemetry.types import TeleportFailureStage
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.teleport.types import (
TeleportAuthRequiredEvent,
TeleportCheckingGitEvent,
TeleportCompleteEvent,
TeleportFetchingUrlEvent,
TeleportPushingEvent,
TeleportPushRequiredEvent,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
TeleportYieldEvent,
)
@ -29,7 +26,6 @@ def send_teleport_early_failure_telemetry(
stage=stage,
error_class=error_class,
push_required=False,
github_auth_required=False,
nb_session_messages=nb_session_messages,
)
@ -40,7 +36,6 @@ class TeleportTelemetryTracker:
nb_session_messages: int
stage: TeleportFailureStage
push_required: bool = False
github_auth_required: bool = False
success: bool = False
error_class: str | None = None
@ -55,13 +50,6 @@ class TeleportTelemetryTracker:
self.stage = "push"
case TeleportStartingWorkflowEvent():
self.stage = "workflow_start"
case TeleportWaitingForGitHubEvent():
self.stage = "github_auth"
case TeleportAuthRequiredEvent():
self.github_auth_required = True
self.stage = "github_auth"
case TeleportFetchingUrlEvent():
self.stage = "fetch_url"
case TeleportCompleteEvent():
self.success = True
@ -78,7 +66,6 @@ class TeleportTelemetryTracker:
def send_success(self) -> None:
self.telemetry_client.send_teleport_completed(
push_required=self.push_required,
github_auth_required=self.github_auth_required,
nb_session_messages=self.nb_session_messages,
)
@ -89,6 +76,5 @@ class TeleportTelemetryTracker:
stage=self.stage,
error_class=self.error_class,
push_required=self.push_required,
github_auth_required=self.github_auth_required,
nb_session_messages=self.nb_session_messages,
)

View file

@ -13,96 +13,61 @@ import zstandard
from vibe.core.config import VibeConfig
from vibe.core.session.session_logger import SessionLogger
from vibe.core.teleport.errors import ServiceTeleportError
from vibe.core.teleport.experimental_nuage import (
ExperimentalNuageClient,
ExperimentalNuageContext,
ExperimentalNuageDiff,
ExperimentalNuageMessage,
ExperimentalNuageRepository,
ExperimentalNuageRequest,
ExperimentalNuageTextPart,
)
from vibe.core.teleport.git import GitRepoInfo, GitRepository
from vibe.core.teleport.nuage import (
ChatAssistantParams,
GitHubParams,
NuageClient,
TeleportSession,
TextChunk,
VibeAgent,
WorkflowConfig,
WorkflowIntegrations,
WorkflowParams,
NuageContext,
NuageDiff,
NuageMessage,
NuageRepository,
NuageRequest,
NuageTextPart,
)
from vibe.core.teleport.types import (
TeleportAuthCompleteEvent,
TeleportAuthRequiredEvent,
TeleportCheckingGitEvent,
TeleportCompleteEvent,
TeleportFetchingUrlEvent,
TeleportPushingEvent,
TeleportPushRequiredEvent,
TeleportPushResponseEvent,
TeleportSendEvent,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
TeleportYieldEvent,
)
from vibe.core.utils.http import build_ssl_context
_DEFAULT_TELEPORT_PROMPT = "Your session has been teleported on a remote workspace. Changes of workspace has been automatically teleported. External workspace changes has NOT been teleported. Environment variables has NOT been teleported. Please continue where you left off."
class TeleportService:
def __init__(
self,
session_logger: SessionLogger,
vibe_code_base_url: str,
vibe_code_workflow_id: str,
vibe_code_sessions_base_url: str,
vibe_code_api_key: str,
workdir: Path | None = None,
*,
vibe_code_task_queue: str | None = None,
vibe_config: VibeConfig | None = None,
client: httpx.AsyncClient | None = None,
timeout: float = 60.0,
) -> None:
self._session_logger = session_logger
self._vibe_code_base_url = vibe_code_base_url
self._vibe_code_workflow_id = vibe_code_workflow_id
self._vibe_code_sessions_base_url = vibe_code_sessions_base_url
self._vibe_code_api_key = vibe_code_api_key
self._vibe_code_task_queue = vibe_code_task_queue
self._vibe_code_project_name = (
vibe_config.vibe_code_project_name if vibe_config else None
)
self._experimental_nuage_enabled = (
vibe_config.vibe_code_experimental_nuage_enabled if vibe_config else False
)
self._vibe_config = vibe_config
self._git = GitRepository(workdir)
self._client = client
self._owns_client = client is None
self._timeout = timeout
self._nuage_client_instance: NuageClient | None = None
self._experimental_nuage_client_instance: ExperimentalNuageClient | None = None
async def __aenter__(self) -> TeleportService:
if self._client is None:
self._client = httpx.AsyncClient(
timeout=httpx.Timeout(self._timeout), verify=build_ssl_context()
)
if self._experimental_nuage_enabled:
self._experimental_nuage_client_instance = ExperimentalNuageClient(
self._vibe_code_base_url, self._vibe_code_api_key, client=self._client
)
else:
self._nuage_client_instance = NuageClient(
self._vibe_code_base_url,
self._vibe_code_api_key,
self._vibe_code_workflow_id,
task_queue=self._vibe_code_task_queue,
client=self._client,
)
self._nuage_client_instance = NuageClient(
self._vibe_code_sessions_base_url,
self._vibe_code_api_key,
client=self._client,
)
await self._git.__aenter__()
return self
@ -130,24 +95,12 @@ class TeleportService:
def _nuage_client(self) -> NuageClient:
if self._nuage_client_instance is None:
self._nuage_client_instance = NuageClient(
self._vibe_code_base_url,
self._vibe_code_sessions_base_url,
self._vibe_code_api_key,
self._vibe_code_workflow_id,
task_queue=self._vibe_code_task_queue,
client=self._http_client,
)
return self._nuage_client_instance
@property
def _experimental_nuage_client(self) -> ExperimentalNuageClient:
if self._experimental_nuage_client_instance is None:
self._experimental_nuage_client_instance = ExperimentalNuageClient(
self._vibe_code_base_url,
self._vibe_code_api_key,
client=self._http_client,
)
return self._experimental_nuage_client_instance
async def check_supported(self) -> None:
await self._git.get_info()
@ -155,23 +108,15 @@ class TeleportService:
return await self._git.is_supported()
async def execute(
self, prompt: str | None, session: TeleportSession
self, prompt: str
) -> AsyncGenerator[TeleportYieldEvent, TeleportSendEvent]:
if prompt:
lechat_user_message = prompt
else:
last_user_message = self._get_last_user_message(session)
if not last_user_message:
raise ServiceTeleportError(
"No prompt provided and no user message found in session."
)
lechat_user_message = f"{last_user_message} (continue)"
prompt = _DEFAULT_TELEPORT_PROMPT
if not prompt:
raise ServiceTeleportError("Teleport requires a non-empty prompt.")
self._validate_config()
git_info = await self._git.get_info()
if self._experimental_nuage_enabled:
self._validate_experimental_nuage_git_info(git_info)
if git_info.branch is None:
raise ServiceTeleportError("Teleport requires a checked-out branch.")
yield TeleportCheckingGitEvent()
await self._git.fetch()
@ -196,63 +141,10 @@ class TeleportService:
yield TeleportStartingWorkflowEvent()
if self._experimental_nuage_enabled:
result = await self._experimental_nuage_client.start(
self._build_experimental_nuage_request(
lechat_user_message=lechat_user_message, git_info=git_info
)
)
yield TeleportCompleteEvent(url=result.url)
return
execution_id = await self._nuage_client.start_workflow(
WorkflowParams(
prompt=prompt,
message=[TextChunk(text=lechat_user_message)],
config=WorkflowConfig(
agent=VibeAgent(
vibe_config=self._vibe_config.model_dump()
if self._vibe_config
else None,
session=session,
)
),
integrations=WorkflowIntegrations(
github=self._build_github_params(git_info),
chat_assistant=ChatAssistantParams(
create_thread=True,
user_message=lechat_user_message,
project_name=self._vibe_code_project_name,
),
),
)
result = await self._nuage_client.start(
self._build_nuage_request(prompt=prompt, git_info=git_info)
)
yield TeleportWaitingForGitHubEvent()
auth_event_sent = False
async for github_data in self._nuage_client.wait_for_github_connection(
execution_id
):
if github_data.connected:
break
if not auth_event_sent and github_data.oauth_url:
yield TeleportAuthRequiredEvent(
oauth_url=github_data.oauth_url, message=github_data.error
)
auth_event_sent = True
if github_data.error:
yield TeleportWaitingForGitHubEvent(message=github_data.error)
yield TeleportAuthCompleteEvent()
yield TeleportFetchingUrlEvent()
chat_url = await self._nuage_client.get_chat_assistant_url(execution_id)
if not chat_url:
raise ServiceTeleportError("Chat assistant URL is not available yet")
yield TeleportCompleteEvent(url=chat_url)
yield TeleportCompleteEvent(url=result.url)
async def _push_or_fail(self) -> None:
if not await self._git.push_current_branch():
@ -267,37 +159,22 @@ class TeleportService:
)
raise ServiceTeleportError(f"{env_var} not set.")
def _build_github_params(self, git_info: GitRepoInfo) -> GitHubParams:
return GitHubParams(
repo=f"{git_info.owner}/{git_info.repo}",
branch=git_info.branch,
commit=git_info.commit,
teleported_diffs=self._compress_diff(git_info.diff or ""),
)
def _build_experimental_nuage_request(
self, *, lechat_user_message: str, git_info: GitRepoInfo
) -> ExperimentalNuageRequest:
if git_info.branch is None:
raise ServiceTeleportError(
"Experimental Nuage teleport requires a checked-out branch."
)
def _build_nuage_request(
self, *, prompt: str, git_info: GitRepoInfo
) -> NuageRequest:
compressed = self._compress_diff(git_info.diff)
diff = (
ExperimentalNuageDiff(content=compressed.decode("ascii"))
NuageDiff(content=compressed.decode("ascii"))
if compressed is not None
else None
)
return ExperimentalNuageRequest(
return NuageRequest(
idempotency_key=str(uuid4()),
message=ExperimentalNuageMessage(
parts=[ExperimentalNuageTextPart(text=lechat_user_message)]
),
context=ExperimentalNuageContext(
message=NuageMessage(parts=[NuageTextPart(text=prompt)]),
context=NuageContext(
repositories=[
ExperimentalNuageRepository(
NuageRepository(
repo_url=git_info.remote_url,
branch=git_info.branch,
commit_sha=git_info.commit,
@ -307,14 +184,6 @@ class TeleportService:
),
)
def _validate_experimental_nuage_git_info(self, git_info: GitRepoInfo) -> None:
if git_info.branch is not None:
return
raise ServiceTeleportError(
"Experimental Nuage teleport requires a checked-out branch."
)
def _compress_diff(self, diff: str, max_size: int = 1_000_000) -> bytes | None:
if not diff:
return None
@ -325,11 +194,3 @@ class TeleportService:
"Diff too large to teleport. Please commit and push your changes first."
)
return encoded
def _get_last_user_message(self, session: TeleportSession) -> str | None:
for msg in reversed(session.messages):
if msg.get("role") == "user":
content = msg.get("content")
if isinstance(content, str) and content:
return content
return None

View file

@ -3,15 +3,6 @@ from __future__ import annotations
from vibe.core.types import BaseEvent
class TeleportAuthRequiredEvent(BaseEvent):
oauth_url: str
message: str | None = None
class TeleportAuthCompleteEvent(BaseEvent):
pass
class TeleportStartingWorkflowEvent(BaseEvent):
pass
@ -33,27 +24,15 @@ class TeleportPushingEvent(BaseEvent):
pass
class TeleportWaitingForGitHubEvent(BaseEvent):
message: str | None = None
class TeleportFetchingUrlEvent(BaseEvent):
pass
class TeleportCompleteEvent(BaseEvent):
url: str
type TeleportYieldEvent = (
TeleportAuthRequiredEvent
| TeleportAuthCompleteEvent
| TeleportCheckingGitEvent
TeleportCheckingGitEvent
| TeleportPushRequiredEvent
| TeleportPushingEvent
| TeleportStartingWorkflowEvent
| TeleportWaitingForGitHubEvent
| TeleportFetchingUrlEvent
| TeleportCompleteEvent
)