feat: introduce Rig agent runtime migration

2026-08-04 02:15:18 -05:00
parent d9cf0d8ae3
commit 4c7270db8d
39 changed files with 2551 additions and 211 deletions
@@ -0,0 +1,11 @@
+//! Provider- and UI-independent contracts for Galaxy agent runtimes.
+//!
+//! This crate is the stable boundary between Galaxy application services and
+//! concrete runtimes such as Rig-backed providers or ACP agents. It must not
+//! depend on GalaxyUI, provider SDKs, persistence, or Warp wire protocols.
+
+mod runtime;
+mod types;
+
+pub use runtime::*;
+pub use types::*;
@@ -0,0 +1,144 @@
+use std::error::Error;
+use std::fmt;
+use std::pin::Pin;
+
+use async_channel::{Receiver, Sender, TrySendError};
+use async_trait::async_trait;
+use futures::Stream;
+use serde::{Deserialize, Serialize};
+
+use crate::{AgentEvent, TurnRequest};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum RuntimeKind {
+    Provider,
+    Acp,
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RuntimeCapabilities {
+    pub model_selection: bool,
+    pub session_resume: bool,
+    pub steering: bool,
+    pub tool_permissions: bool,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RuntimeDescriptor {
+    pub id: String,
+    pub display_name: String,
+    pub kind: RuntimeKind,
+    pub capabilities: RuntimeCapabilities,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum TurnCommand {
+    Cancel,
+    Steer { text: String },
+}
+
+#[derive(Clone, Debug)]
+pub struct TurnCommandSender(Sender<TurnCommand>);
+
+impl TurnCommandSender {
+    pub async fn send(&self, command: TurnCommand) -> Result<(), TurnControlClosed> {
+        self.0.send(command).await.map_err(|_| TurnControlClosed)
+    }
+
+    pub fn try_send(&self, command: TurnCommand) -> Result<(), TrySendError<TurnCommand>> {
+        self.0.try_send(command)
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct TurnControl(Receiver<TurnCommand>);
+
+impl TurnControl {
+    pub async fn receive(&self) -> Result<TurnCommand, TurnControlClosed> {
+        self.0.recv().await.map_err(|_| TurnControlClosed)
+    }
+
+    pub fn try_receive(&self) -> Result<TurnCommand, async_channel::TryRecvError> {
+        self.0.try_recv()
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct TurnControlClosed;
+
+impl fmt::Display for TurnControlClosed {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("turn control channel is closed")
+    }
+}
+
+impl Error for TurnControlClosed {}
+
+pub fn turn_control() -> (TurnCommandSender, TurnControl) {
+    let (sender, receiver) = async_channel::unbounded();
+    (TurnCommandSender(sender), TurnControl(receiver))
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum AgentErrorKind {
+    Configuration,
+    Authentication,
+    RateLimited,
+    ContextWindowExceeded,
+    InvalidRequest,
+    Transport,
+    Provider,
+    Protocol,
+    Tool,
+    Cancelled,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AgentError {
+    pub kind: AgentErrorKind,
+    pub message: String,
+    pub user_message: Option<String>,
+    pub recoverable: bool,
+}
+
+impl AgentError {
+    pub fn new(kind: AgentErrorKind, message: impl Into<String>) -> Self {
+        Self {
+            kind,
+            message: message.into(),
+            user_message: None,
+            recoverable: false,
+        }
+    }
+}
+
+impl fmt::Display for AgentError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(&self.message)
+    }
+}
+
+impl Error for AgentError {}
+
+#[cfg(not(target_family = "wasm"))]
+pub type AgentEventStream =
+    Pin<Box<dyn Stream<Item = Result<AgentEvent, AgentError>> + Send + 'static>>;
+
+#[cfg(target_family = "wasm")]
+pub type AgentEventStream = Pin<Box<dyn Stream<Item = Result<AgentEvent, AgentError>> + 'static>>;
+
+#[cfg_attr(target_family = "wasm", async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait)]
+pub trait AgentRuntime: Send + Sync {
+    fn descriptor(&self) -> &RuntimeDescriptor;
+
+    async fn start_turn(
+        &self,
+        request: TurnRequest,
+        control: TurnControl,
+    ) -> Result<AgentEventStream, AgentError>;
+}
+
+#[cfg(test)]
+#[path = "runtime_tests.rs"]
+mod tests;
@@ -0,0 +1,120 @@
+use std::sync::Mutex;
+
+use futures::{StreamExt, stream};
+
+use super::*;
+use crate::{
+    AgentEvent, ConversationMessage, MessageContent, MessageRole, ModelId, StopReason, Usage,
+};
+
+struct FakeRuntime {
+    descriptor: RuntimeDescriptor,
+    requests: Mutex<Vec<TurnRequest>>,
+    events: Vec<AgentEvent>,
+}
+
+impl FakeRuntime {
+    fn new(events: Vec<AgentEvent>) -> Self {
+        Self {
+            descriptor: RuntimeDescriptor {
+                id: "fake".to_string(),
+                display_name: "Deterministic fake".to_string(),
+                kind: RuntimeKind::Provider,
+                capabilities: RuntimeCapabilities {
+                    model_selection: true,
+                    ..RuntimeCapabilities::default()
+                },
+            },
+            requests: Mutex::new(Vec::new()),
+            events,
+        }
+    }
+}
+
+#[async_trait]
+impl AgentRuntime for FakeRuntime {
+    fn descriptor(&self) -> &RuntimeDescriptor {
+        &self.descriptor
+    }
+
+    async fn start_turn(
+        &self,
+        request: TurnRequest,
+        _control: TurnControl,
+    ) -> Result<AgentEventStream, AgentError> {
+        self.requests.lock().unwrap().push(request);
+        let events = self.events.clone().into_iter().map(Ok);
+        Ok(Box::pin(stream::iter(events)))
+    }
+}
+
+#[test]
+fn fake_runtime_preserves_request_and_event_contract() {
+    futures::executor::block_on(async {
+        let expected_events = vec![
+            AgentEvent::TurnStarted {
+                runtime_request_id: "request-1".to_string(),
+            },
+            AgentEvent::TextDelta {
+                text: "hello".to_string(),
+            },
+            AgentEvent::UsageUpdated {
+                usage: Usage {
+                    input_tokens: 4,
+                    output_tokens: 1,
+                    ..Usage::default()
+                },
+            },
+            AgentEvent::TurnStopped {
+                reason: StopReason::Completed,
+            },
+        ];
+        let fake_runtime = FakeRuntime::new(expected_events.clone());
+        let runtime: &dyn AgentRuntime = &fake_runtime;
+        let request = TurnRequest::new(
+            ModelId::new("fake-model"),
+            vec![ConversationMessage {
+                role: MessageRole::User,
+                content: MessageContent::Text("Say hello".to_string()),
+            }],
+        );
+        let (_sender, control) = turn_control();
+
+        let actual_events = runtime
+            .start_turn(request.clone(), control)
+            .await
+            .unwrap()
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap();
+
+        assert_eq!(runtime.descriptor().id, "fake");
+        assert_eq!(*fake_runtime.requests.lock().unwrap(), vec![request]);
+        assert_eq!(actual_events, expected_events);
+    });
+}
+
+#[test]
+fn turn_control_delivers_cancel_and_steering_in_order() {
+    futures::executor::block_on(async {
+        let (sender, control) = turn_control();
+
+        sender
+            .send(TurnCommand::Steer {
+                text: "focus on tests".to_string(),
+            })
+            .await
+            .unwrap();
+        sender.send(TurnCommand::Cancel).await.unwrap();
+
+        assert_eq!(
+            control.receive().await.unwrap(),
+            TurnCommand::Steer {
+                text: "focus on tests".to_string(),
+            }
+        );
+        assert_eq!(control.receive().await.unwrap(), TurnCommand::Cancel);
+    });
+}
@@ -0,0 +1,227 @@
+use std::collections::BTreeMap;
+
+use serde::{Deserialize, Serialize};
+use serde_json::Value as JsonValue;
+
+pub const MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST: usize = 64_000;
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ConversationMessage {
+    pub role: MessageRole,
+    pub content: MessageContent,
+}
+
+impl ConversationMessage {
+    pub fn truncate_tool_results_for_provider_request(&mut self) {
+        truncate_tool_results_in_content(&mut self.content);
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum MessageRole {
+    User,
+    Assistant,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub enum MessageContent {
+    Text(String),
+    ToolUse {
+        tool_use_id: String,
+        name: String,
+        input: JsonValue,
+    },
+    ToolResult {
+        tool_use_id: String,
+        content: String,
+        is_error: bool,
+    },
+    MultiPart(Vec<ContentPart>),
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub enum ContentPart {
+    Text(String),
+    Image {
+        data: Vec<u8>,
+        mime_type: String,
+    },
+    ToolUse {
+        tool_use_id: String,
+        name: String,
+        input: JsonValue,
+    },
+    ToolResult {
+        tool_use_id: String,
+        content: String,
+        is_error: bool,
+    },
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ToolDefinition {
+    pub name: String,
+    pub description: String,
+    pub input_schema: JsonValue,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct ModelId(String);
+
+impl ModelId {
+    pub fn new(value: impl Into<String>) -> Self {
+        Self(value.into())
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl From<String> for ModelId {
+    fn from(value: String) -> Self {
+        Self(value)
+    }
+}
+
+impl From<&str> for ModelId {
+    fn from(value: &str) -> Self {
+        Self(value.to_string())
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct TurnRequest {
+    pub conversation_id: Option<String>,
+    pub model: ModelId,
+    pub system_prompt: Option<String>,
+    pub messages: Vec<ConversationMessage>,
+    pub tools: Vec<ToolDefinition>,
+    pub max_output_tokens: Option<u64>,
+    pub metadata: BTreeMap<String, JsonValue>,
+}
+
+impl TurnRequest {
+    pub fn new(model: impl Into<ModelId>, messages: Vec<ConversationMessage>) -> Self {
+        Self {
+            conversation_id: None,
+            model: model.into(),
+            system_prompt: None,
+            messages,
+            tools: Vec::new(),
+            max_output_tokens: None,
+            metadata: BTreeMap::new(),
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ToolCall {
+    pub id: String,
+    pub name: String,
+    pub arguments: JsonValue,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct ToolResult {
+    pub call_id: String,
+    pub content: String,
+    pub is_error: bool,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum PermissionKind {
+    Read,
+    Write,
+    Execute,
+    Network,
+    ExternalTool,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct PermissionRequest {
+    pub id: String,
+    pub tool_call: ToolCall,
+    pub kind: PermissionKind,
+    pub reason: Option<String>,
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Usage {
+    pub input_tokens: u64,
+    pub output_tokens: u64,
+    pub cached_input_tokens: u64,
+    pub cache_creation_input_tokens: u64,
+}
+
+impl Usage {
+    pub fn total_tokens(&self) -> u64 {
+        self.input_tokens.saturating_add(self.output_tokens)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum StopReason {
+    Completed,
+    Cancelled,
+    MaxTokens,
+    ContextWindowExceeded,
+    Refusal,
+    ToolLoopLimit,
+    Other(String),
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub enum AgentEvent {
+    TurnStarted { runtime_request_id: String },
+    TextDelta { text: String },
+    ReasoningDelta { text: String },
+    ToolProposed { call: ToolCall },
+    PermissionRequested { request: PermissionRequest },
+    ToolStarted { call: ToolCall },
+    ToolCompleted { result: ToolResult },
+    UsageUpdated { usage: Usage },
+    TurnStopped { reason: StopReason },
+}
+
+fn truncate_tool_results_in_content(content: &mut MessageContent) {
+    match content {
+        MessageContent::Text(_) | MessageContent::ToolUse { .. } => {}
+        MessageContent::ToolResult { content, .. } => truncate_tool_result_text(content),
+        MessageContent::MultiPart(parts) => {
+            for part in parts {
+                if let ContentPart::ToolResult { content, .. } = part {
+                    truncate_tool_result_text(content);
+                }
+            }
+        }
+    }
+}
+
+fn truncate_tool_result_text(content: &mut String) {
+    let char_count = content.chars().count();
+    if char_count <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST {
+        return;
+    }
+
+    let omitted_chars = char_count.saturating_sub(MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST);
+    let marker = format!("\n... [tool result truncated; omitted {omitted_chars} chars] ...\n");
+    let marker_chars = marker.chars().count();
+    let retained_chars = MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST.saturating_sub(marker_chars);
+    let head_chars = retained_chars / 2;
+    let tail_chars = retained_chars.saturating_sub(head_chars);
+    let head: String = content.chars().take(head_chars).collect();
+    let tail: String = content
+        .chars()
+        .rev()
+        .take(tail_chars)
+        .collect::<String>()
+        .chars()
+        .rev()
+        .collect();
+    *content = format!("{head}{marker}{tail}");
+}
+
+#[cfg(test)]
+#[path = "types_tests.rs"]
+mod tests;
@@ -0,0 +1,38 @@
+use super::*;
+
+#[test]
+fn truncates_large_tool_results_for_provider_request() {
+    let prefix = "start:";
+    let suffix = ":end";
+    let middle = "x".repeat(MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST + 1_000);
+    let mut message = ConversationMessage {
+        role: MessageRole::User,
+        content: MessageContent::ToolResult {
+            tool_use_id: "toolu_1".to_string(),
+            content: format!("{prefix}{middle}{suffix}"),
+            is_error: false,
+        },
+    };
+
+    message.truncate_tool_results_for_provider_request();
+
+    let MessageContent::ToolResult { content, .. } = message.content else {
+        panic!("expected tool result");
+    };
+    assert!(content.len() <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST + 128);
+    assert!(content.starts_with(prefix));
+    assert!(content.ends_with(suffix));
+    assert!(content.contains("tool result truncated"));
+}
+
+#[test]
+fn usage_total_excludes_cached_breakdown_to_avoid_double_counting() {
+    let usage = Usage {
+        input_tokens: 100,
+        output_tokens: 25,
+        cached_input_tokens: 80,
+        cache_creation_input_tokens: 10,
+    };
+
+    assert_eq!(usage.total_tokens(), 125);
+}