Initial commit
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Laure Hugo <laure.hugo@mistral.ai> Co-Authored-By: Benjamin Trom <benjamin.trom@mistral.ai> Co-Authored-By: Mathias Gesbert <mathias.gesbert@ext.mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-Authored-By: Valentin Berard <val@mistral.ai> Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Quentin Torroba
Quentin Torroba
commit
fa15fc977b
200 changed files with 30484 additions and 0 deletions
925
tests/acp/test_acp.py
Normal file
925
tests/acp/test_acp.py
Normal file
|
|
@ -0,0 +1,925 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from acp import (
|
||||
InitializeRequest,
|
||||
NewSessionRequest,
|
||||
PromptRequest,
|
||||
ReadTextFileRequest,
|
||||
ReadTextFileResponse,
|
||||
RequestPermissionRequest,
|
||||
RequestPermissionResponse,
|
||||
WriteTextFileRequest,
|
||||
)
|
||||
from acp.schema import (
|
||||
AgentCapabilities,
|
||||
AllowedOutcome,
|
||||
DeniedOutcome,
|
||||
Implementation,
|
||||
InitializeResponse,
|
||||
McpCapabilities,
|
||||
NewSessionResponse,
|
||||
PromptCapabilities,
|
||||
PromptResponse,
|
||||
SessionNotification,
|
||||
TextContentBlock,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
import pytest
|
||||
|
||||
from tests import TESTS_ROOT
|
||||
from tests.mock.utils import get_mocking_env, mock_llm_chunk
|
||||
from vibe.acp.utils import ToolOption
|
||||
from vibe.core.types import FunctionCall, ToolCall
|
||||
|
||||
RESPONSE_TIMEOUT = 2.0
|
||||
MOCK_ENTRYPOINT_PATH = "tests/mock/mock_entrypoint.py"
|
||||
PLAYGROUND_DIR = TESTS_ROOT / "playground"
|
||||
|
||||
|
||||
class JsonRpcRequest(BaseModel):
|
||||
jsonrpc: str = "2.0"
|
||||
id: int | str
|
||||
method: str
|
||||
params: Any | None = None
|
||||
|
||||
|
||||
class JsonRpcError(BaseModel):
|
||||
code: int
|
||||
message: str
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class JsonRpcResponse(BaseModel):
|
||||
jsonrpc: str = "2.0"
|
||||
id: int | str | None = None
|
||||
result: Any | None = None
|
||||
error: JsonRpcError | None = None
|
||||
|
||||
|
||||
class JsonRpcNotification(BaseModel):
|
||||
jsonrpc: str = "2.0"
|
||||
method: str
|
||||
params: Any | None = None
|
||||
|
||||
|
||||
type JsonRpcMessage = JsonRpcResponse | JsonRpcNotification | JsonRpcRequest
|
||||
|
||||
|
||||
class InitializeJsonRpcRequest(JsonRpcRequest):
|
||||
method: str = "initialize"
|
||||
params: InitializeRequest | None = None
|
||||
|
||||
|
||||
class InitializeJsonRpcResponse(JsonRpcResponse):
|
||||
result: InitializeResponse | None = None
|
||||
|
||||
|
||||
class NewSessionJsonRpcRequest(JsonRpcRequest):
|
||||
method: str = "session/new"
|
||||
params: NewSessionRequest | None = None
|
||||
|
||||
|
||||
class NewSessionJsonRpcResponse(JsonRpcResponse):
|
||||
result: NewSessionResponse | None = None
|
||||
|
||||
|
||||
class PromptJsonRpcRequest(JsonRpcRequest):
|
||||
method: str = "session/prompt"
|
||||
params: PromptRequest | None = None
|
||||
|
||||
|
||||
class PromptJsonRpcResponse(JsonRpcResponse):
|
||||
result: PromptResponse | None = None
|
||||
|
||||
|
||||
class UpdateJsonRpcNotification(JsonRpcNotification):
|
||||
method: str = "session/update"
|
||||
params: SessionNotification | None = None
|
||||
|
||||
|
||||
class RequestPermissionJsonRpcRequest(JsonRpcRequest):
|
||||
method: str = "session/request_permission"
|
||||
params: RequestPermissionRequest | None = None
|
||||
|
||||
|
||||
class RequestPermissionJsonRpcResponse(JsonRpcResponse):
|
||||
result: RequestPermissionResponse | None = None
|
||||
|
||||
|
||||
class ReadTextFileJsonRpcRequest(JsonRpcRequest):
|
||||
method: str = "fs/read_text_file"
|
||||
params: ReadTextFileRequest | None = None
|
||||
|
||||
|
||||
class ReadTextFileJsonRpcResponse(JsonRpcResponse):
|
||||
result: ReadTextFileResponse | None = None
|
||||
|
||||
|
||||
class WriteTextFileJsonRpcRequest(JsonRpcRequest):
|
||||
method: str = "fs/write_text_file"
|
||||
params: WriteTextFileRequest | None = None
|
||||
|
||||
|
||||
class WriteTextFileJsonRpcResponse(JsonRpcResponse):
|
||||
result: None = None
|
||||
|
||||
|
||||
async def get_acp_agent_process(
|
||||
mock: bool = True, mock_env: dict[str, str] | None = None
|
||||
) -> AsyncGenerator[asyncio.subprocess.Process]:
|
||||
current_env = os.environ.copy()
|
||||
cmd = ["uv", "run", MOCK_ENTRYPOINT_PATH if mock else "vibe-acp"]
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=TESTS_ROOT.parent,
|
||||
env={
|
||||
**current_env,
|
||||
**(mock_env or {}),
|
||||
**({"MISTRAL_API_KEY": "mock"} if mock else {}),
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
yield process
|
||||
finally:
|
||||
# Cleanup
|
||||
if process.returncode is None:
|
||||
process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=0.5)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
|
||||
|
||||
async def send_json_rpc(
|
||||
process: asyncio.subprocess.Process, message: JsonRpcMessage
|
||||
) -> None:
|
||||
if process.stdin is None:
|
||||
raise RuntimeError("Process stdin not available")
|
||||
|
||||
request = message.model_dump_json()
|
||||
request_json = request + "\n"
|
||||
process.stdin.write(request_json.encode())
|
||||
await process.stdin.drain()
|
||||
|
||||
|
||||
async def read_response(
|
||||
process: asyncio.subprocess.Process, timeout: float = RESPONSE_TIMEOUT
|
||||
) -> str | None:
|
||||
if process.stdout is None:
|
||||
raise RuntimeError("Process stdout not available")
|
||||
|
||||
try:
|
||||
# Keep reading lines until we find a valid JSON line
|
||||
while True:
|
||||
line = await asyncio.wait_for(process.stdout.readline(), timeout=timeout)
|
||||
|
||||
if not line:
|
||||
return None
|
||||
|
||||
line_str = line.decode().strip()
|
||||
if not line_str:
|
||||
continue
|
||||
|
||||
try:
|
||||
json.loads(line_str)
|
||||
return line_str
|
||||
except json.JSONDecodeError:
|
||||
# Not JSON, skip it (it's a log message)
|
||||
continue
|
||||
except TimeoutError:
|
||||
return None
|
||||
|
||||
|
||||
async def read_response_for_id(
|
||||
process: asyncio.subprocess.Process,
|
||||
expected_id: int | str,
|
||||
timeout: float = RESPONSE_TIMEOUT,
|
||||
) -> str | None:
|
||||
loop = asyncio.get_running_loop()
|
||||
end_time = loop.time() + timeout
|
||||
|
||||
while (remaining := end_time - loop.time()) > 0:
|
||||
response = await read_response(process, timeout=remaining)
|
||||
if response is None:
|
||||
return None
|
||||
|
||||
response_json = json.loads(response)
|
||||
if response_json.get("id") == expected_id:
|
||||
return response
|
||||
print(
|
||||
f"Skipping response with id={response_json.get('id')}, expecting {expected_id}"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def read_multiple_responses(
|
||||
process: asyncio.subprocess.Process,
|
||||
max_count: int = 10,
|
||||
timeout_per_response: float = RESPONSE_TIMEOUT,
|
||||
) -> list[str]:
|
||||
responses = []
|
||||
for _ in range(max_count):
|
||||
response = await read_response(process, timeout=timeout_per_response)
|
||||
if response:
|
||||
responses.append(response)
|
||||
else:
|
||||
break
|
||||
return responses
|
||||
|
||||
|
||||
def parse_conversation(message_texts: list[str]) -> list[JsonRpcMessage]:
|
||||
parsed_messages: list[JsonRpcMessage] = []
|
||||
for message_text in message_texts:
|
||||
message_json = json.loads(message_text)
|
||||
cls = None
|
||||
has_method = message_json.get("method", None) is not None
|
||||
has_id = message_json.get("id", None) is not None
|
||||
has_result = message_json.get("result", None) is not None
|
||||
|
||||
is_request = has_method and has_id
|
||||
is_notification = has_method and not has_id
|
||||
is_response = has_result
|
||||
|
||||
if is_request:
|
||||
match message_json.get("method"):
|
||||
case "session/prompt":
|
||||
cls = PromptJsonRpcRequest
|
||||
case "session/request_permission":
|
||||
cls = RequestPermissionJsonRpcRequest
|
||||
case "fs/read_text_file":
|
||||
cls = ReadTextFileJsonRpcRequest
|
||||
case "fs/write_text_file":
|
||||
cls = WriteTextFileJsonRpcRequest
|
||||
elif is_notification:
|
||||
match message_json.get("method"):
|
||||
case "session/update":
|
||||
cls = UpdateJsonRpcNotification
|
||||
elif is_response:
|
||||
# For responses, since we don't know the method, we need to find
|
||||
# the matching request.
|
||||
matching_request = next(
|
||||
(
|
||||
m
|
||||
for m in parsed_messages
|
||||
if isinstance(m, JsonRpcRequest) and m.id == message_json.get("id")
|
||||
),
|
||||
None,
|
||||
)
|
||||
if matching_request is None:
|
||||
# No matching request found in the conversation, it most probably was
|
||||
# not included in the conversation. We use a generic response class.
|
||||
cls = JsonRpcResponse
|
||||
else:
|
||||
match matching_request.method:
|
||||
case "session/prompt":
|
||||
cls = PromptJsonRpcResponse
|
||||
case "session/request_permission":
|
||||
cls = RequestPermissionJsonRpcResponse
|
||||
case "fs/read_text_file":
|
||||
cls = ReadTextFileJsonRpcResponse
|
||||
case "fs/write_text_file":
|
||||
cls = WriteTextFileJsonRpcResponse
|
||||
if cls is None:
|
||||
raise ValueError(f"No valid message class found for {message_json}")
|
||||
parsed_messages.append(cls.model_validate(message_json))
|
||||
return parsed_messages
|
||||
|
||||
|
||||
async def initialize_session(acp_agent_process: asyncio.subprocess.Process) -> str:
|
||||
await send_json_rpc(
|
||||
acp_agent_process,
|
||||
InitializeJsonRpcRequest(id=1, params=InitializeRequest(protocolVersion=1)),
|
||||
)
|
||||
initialize_response = await read_response_for_id(
|
||||
acp_agent_process, expected_id=1, timeout=5.0
|
||||
)
|
||||
assert initialize_response is not None
|
||||
|
||||
await send_json_rpc(
|
||||
acp_agent_process,
|
||||
NewSessionJsonRpcRequest(
|
||||
id=2, params=NewSessionRequest(cwd=str(PLAYGROUND_DIR), mcpServers=[])
|
||||
),
|
||||
)
|
||||
session_response = await read_response_for_id(acp_agent_process, expected_id=2)
|
||||
assert session_response is not None
|
||||
session_response_json = json.loads(session_response)
|
||||
session_response_obj = NewSessionJsonRpcResponse.model_validate(
|
||||
session_response_json
|
||||
)
|
||||
assert session_response_obj.result is not None, "No result in response"
|
||||
return session_response_obj.result.sessionId
|
||||
|
||||
|
||||
class TestInitialization:
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_request_response(self) -> None:
|
||||
mock_env = get_mocking_env()
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
await send_json_rpc(
|
||||
process,
|
||||
InitializeJsonRpcRequest(
|
||||
id=1, params=InitializeRequest(protocolVersion=1)
|
||||
),
|
||||
)
|
||||
|
||||
text_response = await read_response(process, timeout=10.0)
|
||||
assert text_response is not None, "No response to initialize"
|
||||
response_json = json.loads(text_response)
|
||||
response = InitializeJsonRpcResponse.model_validate(response_json)
|
||||
assert response.error is None, f"JSON-RPC error: {response.error}"
|
||||
assert response.result is not None, "No result in response"
|
||||
assert response.result.protocolVersion == 1
|
||||
assert response.result.agentCapabilities == AgentCapabilities(
|
||||
loadSession=False,
|
||||
promptCapabilities=PromptCapabilities(
|
||||
audio=False, embeddedContext=True, image=False
|
||||
),
|
||||
mcpCapabilities=McpCapabilities(http=False, sse=False),
|
||||
)
|
||||
assert response.result.agentInfo == Implementation(
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="0.1.0"
|
||||
)
|
||||
vibe_setup_method = next(
|
||||
(
|
||||
method
|
||||
for method in response.result.authMethods or []
|
||||
if method.id == "vibe-setup"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert vibe_setup_method is not None, "vibe-setup auth not found"
|
||||
assert vibe_setup_method.field_meta is not None
|
||||
assert "terminal-auth" in vibe_setup_method.field_meta.keys()
|
||||
|
||||
|
||||
class TestSessionManagement:
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_sessions_unique_ids(self) -> None:
|
||||
mock_env = get_mocking_env(mock_chunks=[mock_llm_chunk() for _ in range(3)])
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
await send_json_rpc(
|
||||
process,
|
||||
InitializeJsonRpcRequest(
|
||||
id=1, params=InitializeRequest(protocolVersion=1)
|
||||
),
|
||||
)
|
||||
await read_response_for_id(process, expected_id=1, timeout=5.0)
|
||||
|
||||
session_ids = []
|
||||
for i in range(3):
|
||||
await send_json_rpc(
|
||||
process,
|
||||
NewSessionJsonRpcRequest(
|
||||
id=i + 2,
|
||||
params=NewSessionRequest(
|
||||
cwd=str(PLAYGROUND_DIR), mcpServers=[]
|
||||
),
|
||||
),
|
||||
)
|
||||
text_response = await read_response_for_id(
|
||||
process, expected_id=i + 2, timeout=RESPONSE_TIMEOUT
|
||||
)
|
||||
assert text_response is not None
|
||||
response_json = json.loads(text_response)
|
||||
response = NewSessionJsonRpcResponse.model_validate(response_json)
|
||||
assert response.error is None, f"JSON-RPC error: {response.error}"
|
||||
assert response.result is not None, "No result in response"
|
||||
session_ids.append(response.result.sessionId)
|
||||
|
||||
assert len(set(session_ids)) == 3
|
||||
|
||||
|
||||
class TestSessionUpdates:
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_message_chunk_structure(self) -> None:
|
||||
mock_env = get_mocking_env([mock_llm_chunk(content="Hi") for _ in range(2)])
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
# Check stderr for error details if process failed
|
||||
if process.returncode is not None and process.stderr:
|
||||
stderr_data = await process.stderr.read()
|
||||
if stderr_data:
|
||||
# Log stderr for debugging test failures
|
||||
pass # Could add proper logging here if needed
|
||||
|
||||
session_id = await initialize_session(process)
|
||||
|
||||
await send_json_rpc(
|
||||
process,
|
||||
PromptJsonRpcRequest(
|
||||
id=3,
|
||||
params=PromptRequest(
|
||||
sessionId=session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Just say hi")],
|
||||
),
|
||||
),
|
||||
)
|
||||
text_response = await read_response(process)
|
||||
assert text_response is not None
|
||||
response = UpdateJsonRpcNotification.model_validate(
|
||||
json.loads(text_response)
|
||||
)
|
||||
|
||||
assert response.params is not None
|
||||
assert response.params.update.sessionUpdate == "agent_message_chunk"
|
||||
assert response.params.update.content is not None
|
||||
assert response.params.update.content.type == "text"
|
||||
assert response.params.update.content.text is not None
|
||||
assert response.params.update.content.text == "Hi"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_update_structure(self) -> None:
|
||||
mock_env = get_mocking_env([
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="grep", arguments='{"pattern": "auth"}'
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
name="bash",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="The files containing the pattern 'auth' are ...",
|
||||
finish_reason="stop",
|
||||
),
|
||||
])
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
session_id = await initialize_session(process)
|
||||
|
||||
await send_json_rpc(
|
||||
process,
|
||||
PromptJsonRpcRequest(
|
||||
id=3,
|
||||
params=PromptRequest(
|
||||
sessionId=session_id,
|
||||
prompt=[
|
||||
TextContentBlock(
|
||||
type="text",
|
||||
text="Show me files that are related to auth",
|
||||
)
|
||||
],
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=10)
|
||||
assert len(text_responses) > 0
|
||||
responses = [
|
||||
UpdateJsonRpcNotification.model_validate(json.loads(r))
|
||||
for r in text_responses
|
||||
]
|
||||
|
||||
tool_call = next(
|
||||
(
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, UpdateJsonRpcNotification)
|
||||
and r.params is not None
|
||||
and r.params.update.sessionUpdate == "tool_call"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert tool_call is not None
|
||||
assert tool_call.params is not None
|
||||
assert tool_call.params.update is not None
|
||||
|
||||
assert tool_call.params.update.sessionUpdate == "tool_call"
|
||||
assert tool_call.params.update.kind == "search"
|
||||
assert tool_call.params.update.title == "grep: 'auth'"
|
||||
assert (
|
||||
tool_call.params.update.rawInput
|
||||
== '{"pattern":"auth","path":".","max_matches":null,"use_default_ignore":true}'
|
||||
)
|
||||
|
||||
|
||||
async def start_session_with_request_permission(
|
||||
process: asyncio.subprocess.Process, prompt: str
|
||||
) -> RequestPermissionJsonRpcRequest:
|
||||
session_id = await initialize_session(process)
|
||||
await send_json_rpc(
|
||||
process,
|
||||
PromptJsonRpcRequest(
|
||||
id=3,
|
||||
params=PromptRequest(
|
||||
sessionId=session_id,
|
||||
prompt=[TextContentBlock(type="text", text=prompt)],
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(
|
||||
process, max_count=15, timeout_per_response=2.0
|
||||
)
|
||||
|
||||
responses = parse_conversation(text_responses)
|
||||
last_response = responses[-1]
|
||||
|
||||
assert isinstance(last_response, RequestPermissionJsonRpcRequest)
|
||||
assert last_response.params is not None
|
||||
assert len(last_response.params.options) == 3
|
||||
return last_response
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="Disabled until we have a way to properly mock the fs and acp interactions"
|
||||
)
|
||||
class TestToolCallStructure:
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_request_permission_structure(self) -> None:
|
||||
custom_results = [
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="write_file",
|
||||
arguments='{"path":"test.txt","content":"hello, world!"'
|
||||
',"overwrite":true}',
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
name="write_file",
|
||||
finish_reason="stop",
|
||||
),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
session_id = await initialize_session(process)
|
||||
await send_json_rpc(
|
||||
process,
|
||||
PromptJsonRpcRequest(
|
||||
id=3,
|
||||
params=PromptRequest(
|
||||
sessionId=session_id,
|
||||
prompt=[
|
||||
TextContentBlock(
|
||||
type="text",
|
||||
text="Create a new file named test.txt "
|
||||
"with content 'hello, world!'",
|
||||
)
|
||||
],
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=3)
|
||||
responses = parse_conversation(text_responses)
|
||||
|
||||
# Look for tool call request permission updates
|
||||
permission_requests = [
|
||||
r for r in responses if isinstance(r, RequestPermissionJsonRpcRequest)
|
||||
]
|
||||
|
||||
assert len(permission_requests) > 0, (
|
||||
"No tool call permission requests found"
|
||||
)
|
||||
|
||||
first_request = permission_requests[0]
|
||||
assert first_request.params is not None
|
||||
assert first_request.params.toolCall is not None
|
||||
assert first_request.params.toolCall.toolCallId is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_update_approved_structure(self) -> None:
|
||||
custom_results = [
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="write_file",
|
||||
arguments='{"path":"test.txt","content":"hello, world!"'
|
||||
',"overwrite":true}',
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
name="write_file",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="The file test.txt has been created", finish_reason="stop"
|
||||
),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
permission_request = await start_session_with_request_permission(
|
||||
process, "Create a file named test.txt"
|
||||
)
|
||||
assert permission_request.params is not None
|
||||
selected_option_id = ToolOption.ALLOW_ONCE
|
||||
await send_json_rpc(
|
||||
process,
|
||||
RequestPermissionJsonRpcResponse(
|
||||
id=permission_request.id,
|
||||
result=RequestPermissionResponse(
|
||||
outcome=AllowedOutcome(
|
||||
outcome="selected", optionId=selected_option_id
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=7)
|
||||
responses = parse_conversation(text_responses)
|
||||
|
||||
approved_tool_call = next(
|
||||
(
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, UpdateJsonRpcNotification)
|
||||
and r.method == "session/update"
|
||||
and r.params is not None
|
||||
and r.params.update is not None
|
||||
and r.params.update.sessionUpdate == "tool_call_update"
|
||||
and r.params.update.toolCallId
|
||||
== (permission_request.params.toolCall.toolCallId)
|
||||
and r.params.update.status == "completed"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert approved_tool_call is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_update_rejected_structure(self) -> None:
|
||||
custom_results = [
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="write_file",
|
||||
arguments='{"path":"test.txt","content":"hello, world!"'
|
||||
',"overwrite":false}',
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
name="write_file",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="The file test.txt has not been created, "
|
||||
"because you rejected the permission request",
|
||||
finish_reason="stop",
|
||||
),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
permission_request = await start_session_with_request_permission(
|
||||
process, "Create a file named test.txt"
|
||||
)
|
||||
assert permission_request.params is not None
|
||||
selected_option_id = ToolOption.REJECT_ONCE
|
||||
|
||||
await send_json_rpc(
|
||||
process,
|
||||
RequestPermissionJsonRpcResponse(
|
||||
id=permission_request.id,
|
||||
result=RequestPermissionResponse(
|
||||
outcome=AllowedOutcome(
|
||||
outcome="selected", optionId=selected_option_id
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=5)
|
||||
responses = parse_conversation(text_responses)
|
||||
|
||||
rejected_tool_call = next(
|
||||
(
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, UpdateJsonRpcNotification)
|
||||
and r.method == "session/update"
|
||||
and r.params is not None
|
||||
and r.params.update.sessionUpdate == "tool_call_update"
|
||||
and r.params.update.toolCallId
|
||||
== (permission_request.params.toolCall.toolCallId)
|
||||
and r.params.update.status == "failed"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert rejected_tool_call is not None
|
||||
|
||||
@pytest.mark.skip(reason="Long running tool call updates are not implemented yet")
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_in_progress_update_structure(self) -> None:
|
||||
custom_results = [
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="bash",
|
||||
arguments='{"command":"sleep 3","timeout":null}',
|
||||
),
|
||||
type="function",
|
||||
)
|
||||
],
|
||||
name="bash",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="The command sleep 3 has been run", finish_reason="stop"
|
||||
),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
session_id = await initialize_session(process)
|
||||
await send_json_rpc(
|
||||
process,
|
||||
PromptJsonRpcRequest(
|
||||
id=3,
|
||||
params=PromptRequest(
|
||||
sessionId=session_id,
|
||||
prompt=[
|
||||
TextContentBlock(
|
||||
type="text", text="Run sleep 3 in the current directory"
|
||||
)
|
||||
],
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=4)
|
||||
responses = parse_conversation(text_responses)
|
||||
|
||||
# Look for tool call in progress updates
|
||||
in_progress_calls = [
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, UpdateJsonRpcNotification)
|
||||
and r.params is not None
|
||||
and r.params.update.sessionUpdate == "tool_call_update"
|
||||
and r.params.update.status == "in_progress"
|
||||
]
|
||||
|
||||
assert len(in_progress_calls) > 0, (
|
||||
"No tool call in progress updates found for a long running command"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_result_update_failure_structure(self) -> None:
|
||||
custom_results = [
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="write_file",
|
||||
arguments='{"path":"/test.txt","content":"hello, world!"'
|
||||
',"overwrite":true}',
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
name="write_file",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="The file /test.txt has not been created "
|
||||
"because it's outside the project directory",
|
||||
finish_reason="stop",
|
||||
),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
permission_request = await start_session_with_request_permission(
|
||||
process, "Create a file named /test.txt"
|
||||
)
|
||||
assert permission_request.params is not None
|
||||
selected_option_id = ToolOption.ALLOW_ONCE
|
||||
await send_json_rpc(
|
||||
process,
|
||||
RequestPermissionJsonRpcResponse(
|
||||
id=permission_request.id,
|
||||
result=RequestPermissionResponse(
|
||||
outcome=AllowedOutcome(
|
||||
outcome="selected", optionId=selected_option_id
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=7)
|
||||
responses = parse_conversation(text_responses)
|
||||
|
||||
# Look for tool call result failure updates
|
||||
failure_result = next(
|
||||
(
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, UpdateJsonRpcNotification)
|
||||
and r.params is not None
|
||||
and r.params.update.sessionUpdate == "tool_call_update"
|
||||
and r.params.update.status == "failed"
|
||||
and r.params.update.rawOutput is not None
|
||||
and r.params.update.toolCallId is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert failure_result is not None
|
||||
|
||||
|
||||
class TestCancellationStructure:
|
||||
@pytest.mark.skip(
|
||||
reason="Proper cancellation is not implemented yet, we still need to return "
|
||||
"the right end_turn and be able to cancel at any point in time "
|
||||
"(and not only at tool call time)"
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_update_cancelled_structure(self) -> None:
|
||||
custom_results = [
|
||||
mock_llm_chunk(content="Hey"),
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="write_file",
|
||||
arguments='{"path":"test.txt","content":"hello, world!"'
|
||||
',"overwrite":false}',
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
name="write_file",
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="The file test.txt has not been created, "
|
||||
"because you cancelled the permission request",
|
||||
finish_reason="stop",
|
||||
),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
permission_request = await start_session_with_request_permission(
|
||||
process, "Create a file named test.txt"
|
||||
)
|
||||
assert permission_request.params is not None
|
||||
|
||||
await send_json_rpc(
|
||||
process,
|
||||
RequestPermissionJsonRpcResponse(
|
||||
id=permission_request.id,
|
||||
result=RequestPermissionResponse(
|
||||
outcome=DeniedOutcome(outcome="cancelled")
|
||||
),
|
||||
),
|
||||
)
|
||||
text_responses = await read_multiple_responses(process, max_count=5)
|
||||
responses = parse_conversation(text_responses)
|
||||
|
||||
assert len(responses) == 2, (
|
||||
"There should be only 2 responses: "
|
||||
"the tool call update and the prompt end turn"
|
||||
)
|
||||
|
||||
cancelled_tool_call = next(
|
||||
(
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, UpdateJsonRpcNotification)
|
||||
and r.method == "session/update"
|
||||
and r.params is not None
|
||||
and r.params.update.sessionUpdate == "tool_call_update"
|
||||
and r.params.update.toolCallId
|
||||
== (permission_request.params.toolCall.toolCallId)
|
||||
and r.params.update.status == "failed"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert cancelled_tool_call is not None
|
||||
|
||||
cancelled_prompt_response = next(
|
||||
(
|
||||
r
|
||||
for r in responses
|
||||
if isinstance(r, PromptJsonRpcResponse)
|
||||
and r.result is not None
|
||||
and r.result.stopReason == "cancelled"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert cancelled_prompt_response is not None
|
||||
Loading…
Add table
Add a link
Reference in a new issue