v2.12.0 (#725)
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:
parent
cf3f4ca58f
commit
1843196d88
20 changed files with 399 additions and 1675 deletions
|
|
@ -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,
|
||||
},
|
||||
|
|
|
|||
45
tests/core/test_last_user_message.py
Normal file
45
tests/core/test_last_user_message.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
@ -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
|
||||
|
|
@ -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 test_github_params_with_values(self) -> None:
|
||||
params = GitHubParams(
|
||||
repo="owner/repo",
|
||||
branch="main",
|
||||
commit="abc123",
|
||||
pr_number=42,
|
||||
teleported_diffs=b"base64data",
|
||||
)
|
||||
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"
|
||||
)
|
||||
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"),
|
||||
)
|
||||
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",
|
||||
branch="main",
|
||||
commit="abc123",
|
||||
pr_number=42,
|
||||
teleported_diffs=b"base64data",
|
||||
),
|
||||
chat_assistant=ChatAssistantParams(user_message="test"),
|
||||
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")
|
||||
]
|
||||
),
|
||||
)
|
||||
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"
|
||||
|
||||
|
||||
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_posts_nuage_request() -> None:
|
||||
seen_request: httpx.Request | None = None
|
||||
|
||||
@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()
|
||||
|
||||
|
||||
class TestNuageClientStartWorkflow:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def nuage(self, mock_client: MagicMock) -> NuageClient:
|
||||
return NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=mock_client
|
||||
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",
|
||||
},
|
||||
)
|
||||
|
||||
@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)
|
||||
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())
|
||||
|
||||
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)
|
||||
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"
|
||||
|
||||
|
||||
class TestNuageClientGetGitHubIntegration:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_omits_empty_branch() -> None:
|
||||
seen_body: dict[str, object] | None = None
|
||||
|
||||
@pytest.fixture
|
||||
def nuage(self, mock_client: MagicMock) -> NuageClient:
|
||||
return NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=mock_client
|
||||
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",
|
||||
},
|
||||
)
|
||||
|
||||
@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}
|
||||
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"}]},
|
||||
}
|
||||
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.mark.asyncio
|
||||
async def test_start_serializes_commit_sha_and_diff() -> None:
|
||||
seen_body: dict[str, object] | None = None
|
||||
|
||||
@pytest.fixture
|
||||
def nuage(self, mock_client: MagicMock) -> NuageClient:
|
||||
return NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=mock_client
|
||||
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",
|
||||
},
|
||||
)
|
||||
|
||||
@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"}
|
||||
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_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=",
|
||||
}
|
||||
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"
|
||||
@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:
|
||||
nuage = NuageClient("https://chat.example.com", "api-key", client=client)
|
||||
with pytest.raises(ServiceTeleportError, match="Nuage start"):
|
||||
await 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:
|
||||
nuage = NuageClient("https://chat.example.com", "api-key", client=client)
|
||||
with pytest.raises(ServiceTeleportError, match="response was invalid"):
|
||||
await nuage.start(_request())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,94 +13,59 @@ 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_sessions_base_url,
|
||||
self._vibe_code_api_key,
|
||||
self._vibe_code_workflow_id,
|
||||
task_queue=self._vibe_code_task_queue,
|
||||
client=self._client,
|
||||
)
|
||||
await self._git.__aenter__()
|
||||
|
|
@ -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
|
||||
)
|
||||
result = await self._nuage_client.start(
|
||||
self._build_nuage_request(prompt=prompt, 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,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue