tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

draft: use rust messages in typescript (#1393)

Browse Source
This commit is contained in:
Bradley Axen
2025-02-27 04:02:43 +01:00
committed by GitHub
parent 552facb7ef
commit 0602b35ddc
19 changed files with 1224 additions and 1537 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/reply.rs
+78 -308
View File
@@ -10,9 +10,8 @@ use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::message::{Message, MessageContent};
use mcp_core::{content::Content, role::Role};
use serde::Deserialize;
use serde_json::{json, Value};
use mcp_core::role::Role;
use serde::{Deserialize, Serialize};
use std::{
convert::Infallible,
pin::Pin,
@@ -23,33 +22,13 @@ use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_stream::wrappers::ReceiverStream;
// Types matching the incoming JSON structure
// Direct message serialization for the chat request
#[derive(Debug, Deserialize)]
struct ChatRequest {
messages: Vec<IncomingMessage>,
messages: Vec<Message>,
}
#[derive(Debug, Deserialize)]
struct IncomingMessage {
role: String,
content: String,
#[serde(default)]
#[serde(rename = "toolInvocations")]
tool_invocations: Vec<ToolInvocation>,
}
#[derive(Debug, Deserialize)]
struct ToolInvocation {
state: String,
#[serde(rename = "toolCallId")]
tool_call_id: String,
#[serde(rename = "toolName")]
tool_name: String,
args: Value,
result: Option<Vec<Content>>,
}
// Custom SSE response type that implements the Vercel AI SDK protocol
// Custom SSE response type for streaming messages
pub struct SseResponse {
rx: ReceiverStream<String>,
}
@@ -79,188 +58,32 @@ impl IntoResponse for SseResponse {
.header("Content-Type", "text/event-stream")
.header("Cache-Control", "no-cache")
.header("Connection", "keep-alive")
.header("x-vercel-ai-data-stream", "v1")
.body(body)
.unwrap()
}
}
// Convert incoming messages to our internal Message type
fn convert_messages(incoming: Vec<IncomingMessage>) -> Vec<Message> {
let mut messages = Vec::new();
for msg in incoming {
match msg.role.as_str() {
"user" => {
messages.push(Message::user().with_text(msg.content));
}
"assistant" => {
// First handle any tool invocations - each represents a complete request/response cycle
for tool in msg.tool_invocations {
if tool.state == "result" {
// Add the original tool request from assistant
let tool_call = mcp_core::tool::ToolCall {
name: tool.tool_name,
arguments: tool.args,
};
messages.push(
Message::assistant()
.with_tool_request(tool.tool_call_id.clone(), Ok(tool_call)),
);
// Add the tool response from user
if let Some(result) = &tool.result {
messages.push(
Message::user()
.with_tool_response(tool.tool_call_id, Ok(result.clone())),
);
}
}
}
// Then add the assistant's text response after tool interactions
if !msg.content.is_empty() {
messages.push(Message::assistant().with_text(msg.content));
}
}
_ => {
tracing::warn!("Unknown role: {}", msg.role);
}
}
}
messages
// Message event types for SSE streaming
#[derive(Debug, Serialize)]
#[serde(tag = "type")]
enum MessageEvent {
Message { message: Message },
Error { error: String },
Finish { reason: String },
}
// Protocol-specific message formatting
struct ProtocolFormatter;
impl ProtocolFormatter {
fn format_text(text: &str) -> String {
let encoded_text = serde_json::to_string(text).unwrap_or_else(|_| String::new());
format!("0:{}\n", encoded_text)
}
fn format_tool_call(id: &str, name: &str, args: &Value) -> String {
// Tool calls start with "9:"
let tool_call = json!({
"toolCallId": id,
"toolName": name,
"args": args
});
format!("9:{}\n", tool_call)
}
fn format_tool_response(id: &str, result: &Vec<Content>) -> String {
// Tool responses start with "a:"
let response = json!({
"toolCallId": id,
"result": result,
});
format!("a:{}\n", response)
}
fn format_error(error: &str) -> String {
// Error messages start with "3:" in the new protocol.
let encoded_error = serde_json::to_string(error).unwrap_or_else(|_| String::new());
format!("3:{}\n", encoded_error)
}
fn format_finish(reason: &str) -> String {
// Finish messages start with "d:"
let finish = json!({
"finishReason": reason,
"usage": {
"promptTokens": 0,
"completionTokens": 0
}
});
format!("d:{}\n", finish)
}
}
async fn stream_message(
message: Message,
// Stream a message as an SSE event
async fn stream_event(
event: MessageEvent,
tx: &mpsc::Sender<String>,
) -> Result<(), mpsc::error::SendError<String>> {
match message.role {
Role::User => {
// Handle tool responses
for content in message.content {
// I believe with the protocol we aren't intended to pass back user messages, so we only deal with
// the tool responses here
if let MessageContent::ToolResponse(response) = content {
// We should return a result for either an error or a success
match response.tool_result {
Ok(result) => {
tx.send(ProtocolFormatter::format_tool_response(
&response.id,
&result,
))
.await?;
}
Err(err) => {
// Send an error message first
tx.send(ProtocolFormatter::format_error(&err.to_string()))
.await?;
// Then send an empty tool response to maintain the protocol
let result =
vec![Content::text(format!("Error: {}", err)).with_priority(0.0)];
tx.send(ProtocolFormatter::format_tool_response(
&response.id,
&result,
))
.await?;
}
}
}
}
}
Role::Assistant => {
for content in message.content {
match content {
MessageContent::ToolRequest(request) => {
match request.tool_call {
Ok(tool_call) => {
tx.send(ProtocolFormatter::format_tool_call(
&request.id,
&tool_call.name,
&tool_call.arguments,
))
.await?;
}
Err(err) => {
// Send a placeholder tool call to maintain protocol
tx.send(ProtocolFormatter::format_tool_call(
&request.id,
"invalid_tool",
&json!({"error": err.to_string()}),
))
.await?;
}
}
}
MessageContent::Text(text) => {
for line in text.text.lines() {
let modified_line = format!("{}\n", line);
tx.send(ProtocolFormatter::format_text(&modified_line))
.await?;
}
}
MessageContent::ToolConfirmationRequest(_) => {
// skip tool confirmation requests
}
MessageContent::Image(_) => {
// skip images
}
MessageContent::ToolResponse(_) => {
// skip tool responses
}
}
}
}
}
Ok(())
let json = serde_json::to_string(&event).unwrap_or_else(|e| {
format!(
r#"{{"type":"Error","error":"Failed to serialize event: {}"}}"#,
e
)
});
tx.send(format!("data: {}\n\n", json)).await
}
async fn handler(
@@ -278,19 +101,12 @@ async fn handler(
return Err(StatusCode::UNAUTHORIZED);
}
// Check protocol header (optional in our case)
if let Some(protocol) = headers.get("x-protocol") {
if protocol.to_str().map(|p| p != "data").unwrap_or(true) {
return Err(StatusCode::BAD_REQUEST);
}
}
// Create channel for streaming
let (tx, rx) = mpsc::channel(100);
let stream = ReceiverStream::new(rx);
// Convert incoming messages
let messages = convert_messages(request.messages);
// Get messages directly from the request
let messages = request.messages;
// Get a lock on the shared agent
let agent = state.agent.clone();
@@ -301,10 +117,20 @@ async fn handler(
let agent = match agent.as_ref() {
Some(agent) => agent,
None => {
let _ = tx
.send(ProtocolFormatter::format_error("No agent configured"))
.await;
let _ = tx.send(ProtocolFormatter::format_finish("error")).await;
let _ = stream_event(
MessageEvent::Error {
error: "No agent configured".to_string(),
},
&tx,
)
.await;
let _ = stream_event(
MessageEvent::Finish {
reason: "error".to_string(),
},
&tx,
)
.await;
return;
}
};
@@ -313,10 +139,20 @@ async fn handler(
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
let _ = tx
.send(ProtocolFormatter::format_error(&e.to_string()))
.await;
let _ = tx.send(ProtocolFormatter::format_finish("error")).await;
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
)
.await;
let _ = stream_event(
MessageEvent::Finish {
reason: "error".to_string(),
},
&tx,
)
.await;
return;
}
};
@@ -326,25 +162,32 @@ async fn handler(
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(message))) => {
if let Err(e) = stream_message(message, &tx).await {
if let Err(e) = stream_event(MessageEvent::Message { message }, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await;
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
break;
}
}
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await;
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
break;
}
Ok(None) => {
break;
}
Err(_) => { // Heartbeat, used to detect disconnected clients and then end running tools.
Err(_) => { // Heartbeat, used to detect disconnected clients
if tx.is_closed() {
// Kill any running processes when the client disconnects
// TODO is this used? I suspect post MCP this is on the server instead
// goose::process_store::kill_processes();
break;
}
continue;
@@ -354,24 +197,30 @@ async fn handler(
}
}
// Send finish message
let _ = tx.send(ProtocolFormatter::format_finish("stop")).await;
// Send finish event
let _ = stream_event(
MessageEvent::Finish {
reason: "stop".to_string(),
},
&tx,
)
.await;
});
Ok(SseResponse::new(stream))
}
#[derive(Debug, Deserialize, serde::Serialize)]
#[derive(Debug, Deserialize, Serialize)]
struct AskRequest {
prompt: String,
}
#[derive(Debug, serde::Serialize)]
#[derive(Debug, Serialize)]
struct AskResponse {
response: String,
}
// simple ask an AI for a response, non streaming
// Simple ask an AI for a response, non streaming
async fn ask_handler(
State(state): State<AppState>,
headers: HeaderMap,
@@ -478,85 +327,6 @@ mod tests {
}
}
#[test]
fn test_convert_messages_user_only() {
let incoming = vec![IncomingMessage {
role: "user".to_string(),
content: "Hello".to_string(),
tool_invocations: vec![],
}];
let messages = convert_messages(incoming);
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::User);
assert!(
matches!(&messages[0].content[0], MessageContent::Text(text) if text.text == "Hello")
);
}
#[test]
fn test_convert_messages_with_tool_invocation() {
let tool_result = vec![Content::text("tool response").with_priority(0.0)];
let incoming = vec![IncomingMessage {
role: "assistant".to_string(),
content: "".to_string(),
tool_invocations: vec![ToolInvocation {
state: "result".to_string(),
tool_call_id: "123".to_string(),
tool_name: "test_tool".to_string(),
args: json!({"key": "value"}),
result: Some(tool_result.clone()),
}],
}];
let messages = convert_messages(incoming);
assert_eq!(messages.len(), 2); // Tool request and response
// Check tool request
assert_eq!(messages[0].role, Role::Assistant);
assert!(
matches!(&messages[0].content[0], MessageContent::ToolRequest(req) if req.id == "123")
);
// Check tool response
assert_eq!(messages[1].role, Role::User);
assert!(
matches!(&messages[1].content[0], MessageContent::ToolResponse(resp) if resp.id == "123")
);
}
#[test]
fn test_protocol_formatter() {
// Test text formatting
let text = "Hello world";
let formatted = ProtocolFormatter::format_text(text);
assert_eq!(formatted, "0:\"Hello world\"\n");
// Test tool call formatting
let formatted =
ProtocolFormatter::format_tool_call("123", "test_tool", &json!({"key": "value"}));
assert!(formatted.starts_with("9:"));
assert!(formatted.contains("\"toolCallId\":\"123\""));
assert!(formatted.contains("\"toolName\":\"test_tool\""));
// Test tool response formatting
let result = vec![Content::text("response").with_priority(0.0)];
let formatted = ProtocolFormatter::format_tool_response("123", &result);
assert!(formatted.starts_with("a:"));
assert!(formatted.contains("\"toolCallId\":\"123\""));
// Test error formatting
let formatted = ProtocolFormatter::format_error("Test error");
println!("Formatted error: {}", formatted);
assert!(formatted.starts_with("3:"));
assert!(formatted.contains("Test error"));
// Test finish formatting
let formatted = ProtocolFormatter::format_finish("stop");
assert!(formatted.starts_with("d:"));
assert!(formatted.contains("\"finishReason\":\"stop\""));
}
mod integration_tests {
use super::*;
use axum::{body::Body, http::Request};
@@ -575,7 +345,7 @@ mod tests {
});
let agent = AgentFactory::create("reference", mock_provider).unwrap();
let state = AppState {
config: Arc::new(Mutex::new(HashMap::new())), // Add this line
config: Arc::new(Mutex::new(HashMap::new())),
agent: Arc::new(Mutex::new(Some(agent))),
secret_key: "test-secret".to_string(),
};
crates/goose/src/message.rs
+129
View File
@@ -14,19 +14,26 @@ use mcp_core::role::Role;
use mcp_core::tool::ToolCall;
use serde_json::Value;
mod tool_result_serde;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolRequest {
pub id: String,
#[serde(with = "tool_result_serde")]
pub tool_call: ToolResult<ToolCall>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolResponse {
pub id: String,
#[serde(with = "tool_result_serde")]
pub tool_result: ToolResult<Vec<Content>>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolConfirmationRequest {
pub id: String,
pub tool_name: String,
@@ -36,6 +43,7 @@ pub struct ToolConfirmationRequest {
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
/// Content passed inside a message, which can be both simple content and tool content
#[serde(tag = "type", rename_all = "camelCase")]
pub enum MessageContent {
Text(TextContent),
Image(ImageContent),
@@ -150,6 +158,7 @@ impl From<Content> for MessageContent {
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
/// A message to or from an LLM
#[serde(rename_all = "camelCase")]
pub struct Message {
pub role: Role,
pub created: i64,
@@ -292,3 +301,123 @@ impl Message {
.all(|c| matches!(c, MessageContent::Text(_)))
}
}
#[cfg(test)]
mod tests {
use super::*;
use mcp_core::handler::ToolError;
use serde_json::{json, Value};
#[test]
fn test_message_serialization() {
let message = Message::assistant()
.with_text("Hello, I'll help you with that.")
.with_tool_request(
"tool123",
Ok(ToolCall::new("test_tool", json!({"param": "value"}))),
);
let json_str = serde_json::to_string_pretty(&message).unwrap();
println!("Serialized message: {}", json_str);
// Parse back to Value to check structure
let value: Value = serde_json::from_str(&json_str).unwrap();
// Check top-level fields
assert_eq!(value["role"], "assistant");
assert!(value["created"].is_i64());
assert!(value["content"].is_array());
// Check content items
let content = &value["content"];
// First item should be text
assert_eq!(content[0]["type"], "text");
assert_eq!(content[0]["text"], "Hello, I'll help you with that.");
// Second item should be toolRequest
assert_eq!(content[1]["type"], "toolRequest");
assert_eq!(content[1]["id"], "tool123");
// Check tool_call serialization
assert_eq!(content[1]["toolCall"]["status"], "success");
assert_eq!(content[1]["toolCall"]["value"]["name"], "test_tool");
assert_eq!(
content[1]["toolCall"]["value"]["arguments"]["param"],
"value"
);
}
#[test]
fn test_error_serialization() {
let message = Message::assistant().with_tool_request(
"tool123",
Err(ToolError::ExecutionError(
"Something went wrong".to_string(),
)),
);
let json_str = serde_json::to_string_pretty(&message).unwrap();
println!("Serialized error: {}", json_str);
// Parse back to Value to check structure
let value: Value = serde_json::from_str(&json_str).unwrap();
// Check tool_call serialization with error
let tool_call = &value["content"][0]["toolCall"];
assert_eq!(tool_call["status"], "error");
assert_eq!(tool_call["error"], "Execution failed: Something went wrong");
}
#[test]
fn test_deserialization() {
// Create a JSON string with our new format
let json_str = r#"{
"role": "assistant",
"created": 1740171566,
"content": [
{
"type": "text",
"text": "I'll help you with that."
},
{
"type": "toolRequest",
"id": "tool123",
"toolCall": {
"status": "success",
"value": {
"name": "test_tool",
"arguments": {"param": "value"}
}
}
}
]
}"#;
let message: Message = serde_json::from_str(json_str).unwrap();
assert_eq!(message.role, Role::Assistant);
assert_eq!(message.created, 1740171566);
assert_eq!(message.content.len(), 2);
// Check first content item
if let MessageContent::Text(text) = &message.content[0] {
assert_eq!(text.text, "I'll help you with that.");
} else {
panic!("Expected Text content");
}
// Check second content item
if let MessageContent::ToolRequest(req) = &message.content[1] {
assert_eq!(req.id, "tool123");
if let Ok(tool_call) = &req.tool_call {
assert_eq!(tool_call.name, "test_tool");
assert_eq!(tool_call.arguments, json!({"param": "value"}));
} else {
panic!("Expected successful tool call");
}
} else {
panic!("Expected ToolRequest content");
}
}
}
crates/goose/src/message/tool_result_serde.rs
+64
View File
@@ -0,0 +1,64 @@
use mcp_core::handler::{ToolError, ToolResult};
use serde::ser::SerializeStruct;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub fn serialize<T, S>(value: &ToolResult<T>, serializer: S) -> Result<S::Ok, S::Error>
where
T: Serialize,
S: Serializer,
{
match value {
Ok(val) => {
let mut state = serializer.serialize_struct("ToolResult", 2)?;
state.serialize_field("status", "success")?;
state.serialize_field("value", val)?;
state.end()
}
Err(err) => {
let mut state = serializer.serialize_struct("ToolResult", 2)?;
state.serialize_field("status", "error")?;
state.serialize_field("error", &err.to_string())?;
state.end()
}
}
}
// For deserialization, let's use a simpler approach that works with the format we're serializing to
pub fn deserialize<'de, T, D>(deserializer: D) -> Result<ToolResult<T>, D::Error>
where
T: Deserialize<'de>,
D: Deserializer<'de>,
{
// Define a helper enum to handle the two possible formats
#[derive(Deserialize)]
#[serde(untagged)]
enum ResultFormat<T> {
Success { status: String, value: T },
Error { status: String, error: String },
}
let format = ResultFormat::deserialize(deserializer)?;
match format {
ResultFormat::Success { status, value } => {
if status == "success" {
Ok(Ok(value))
} else {
Err(serde::de::Error::custom(format!(
"Expected status 'success', got '{}'",
status
)))
}
}
ResultFormat::Error { status, error } => {
if status == "error" {
Ok(Err(ToolError::ExecutionError(error)))
} else {
Err(serde::de::Error::custom(format!(
"Expected status 'error', got '{}'",
status
)))
}
}
}
}