Get tests working

Co-authored-by: Antonio Scandurra <me@as-cii.com>

crates/agent2/src/thread.rs

@@ -920,10 +920,7 @@ pub struct ToolCallEventStream {
 
 impl ToolCallEventStream {
     #[cfg(test)]
-    pub fn test() -> (
-        Self,
-        mpsc::UnboundedReceiver<Result<AgentResponseEvent, LanguageModelCompletionError>>,
-    ) {
+    pub fn test() -> (Self, ToolCallEventStreamReceiver) {
         let (events_tx, events_rx) =
             mpsc::unbounded::<Result<AgentResponseEvent, LanguageModelCompletionError>>();
 
@@ -939,7 +936,7 @@ impl ToolCallEventStream {
             AgentResponseEventStream(events_tx),
         );
 
-        (stream, events_rx)
+        (stream, ToolCallEventStreamReceiver(events_rx))
     }
 
     fn new(
@@ -975,3 +972,36 @@ impl ToolCallEventStream {
         )
     }
 }
+
+#[cfg(test)]
+pub struct ToolCallEventStreamReceiver(
+    mpsc::UnboundedReceiver<Result<AgentResponseEvent, LanguageModelCompletionError>>,
+);
+
+#[cfg(test)]
+impl ToolCallEventStreamReceiver {
+    pub async fn expect_tool_authorization(&mut self) -> ToolCallAuthorization {
+        let event = self.0.next().await;
+        if let Some(Ok(AgentResponseEvent::ToolCallAuthorization(auth))) = event {
+            auth
+        } else {
+            panic!("Expected ToolCallAuthorization but got: {:?}", event);
+        }
+    }
+}
+
+#[cfg(test)]
+impl std::ops::Deref for ToolCallEventStreamReceiver {
+    type Target = mpsc::UnboundedReceiver<Result<AgentResponseEvent, LanguageModelCompletionError>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+#[cfg(test)]
+impl std::ops::DerefMut for ToolCallEventStreamReceiver {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

crates/assistant_tools/src/edit_agent.rs

@@ -29,7 +29,6 @@ use serde::{Deserialize, Serialize};
 use std::{cmp, iter, mem, ops::Range, path::PathBuf, pin::Pin, sync::Arc, task::Poll};
 use streaming_diff::{CharOperation, StreamingDiff};
 use streaming_fuzzy_matcher::StreamingFuzzyMatcher;
-use util::debug_panic;
 
 #[derive(Serialize)]
 struct CreateFilePromptTemplate {
@@ -682,11 +681,6 @@ impl EditAgent {
                 if last_message.content.is_empty() {
                     conversation.messages.pop();
                 }
-            } else {
-                debug_panic!(
-                    "Last message must be an Assistant tool calling! Got {:?}",
-                    last_message.content
-                );
             }
         }