refactor(tools): add context to report_issue tool (#980)

- Add the ability to add context to `Tool`s, without passing in awkward args to the generic invoke function.
- Graceful tool error if context fails to be set (should be impossible).

Author: hayemaxi

crates/q_cli/src/cli/chat/mod.rs

@@ -1013,13 +1013,7 @@ where
                 style::Print(format!("{}\n", "▔".repeat(terminal_width))),
                 style::SetForegroundColor(Color::Reset),
             )?;
-            let invoke_result = tool
-                .1
-                .invoke(&self.ctx, &mut self.output, GhIssueContext {
-                    conversation_state: &self.conversation_state,
-                    failed_request_ids: &self.failed_request_ids,
-                })
-                .await;
+            let invoke_result = tool.1.invoke(&self.ctx, &mut self.output).await;
 
             if self.interactive && self.spinner.is_some() {
                 queue!(
@@ -1317,6 +1311,9 @@ where
                 .utterance_id(self.conversation_state.message_id().map(|s| s.to_string()));
             match Tool::try_from(tool_use) {
                 Ok(mut tool) => {
+                    // Apply non-Q-generated context to tools
+                    self.contextualize_tool(&mut tool);
+
                     match tool.validate(&self.ctx).await {
                         Ok(()) => {
                             tool_telemetry.is_valid = Some(true);
@@ -1400,6 +1397,27 @@ where
         }
     }
 
+    /// Apply program context to tools that Q may not have.
+    // We cannot attach this any other way because Tools are constructed by deserializing
+    // output from Amazon Q.
+    // TODO: Is there a better way?
+    fn contextualize_tool(&self, tool: &mut Tool) {
+        #[allow(clippy::single_match)]
+        match tool {
+            Tool::GhIssue(gh_issue) => {
+                gh_issue.set_context(GhIssueContext {
+                    // Ideally we avoid cloning, but this function is not called very often.
+                    // Using references with lifetimes requires a large refactor, and Arc<Mutex<T>>
+                    // seems like overkill and may incur some performance cost anyway.
+                    context_manager: self.conversation_state.context_manager.clone(),
+                    transcript: self.conversation_state.transcript.clone(),
+                    failed_request_ids: self.failed_request_ids.clone(),
+                });
+            },
+            _ => (),
+        };
+    }
+
     async fn print_tool_descriptions(&mut self, tool_uses: &[QueuedTool]) -> Result<(), ChatError> {
         let terminal_width = self.terminal_width();
         for (_, tool) in tool_uses.iter() {

crates/q_cli/src/cli/chat/tools/gh_issue.rs

@@ -1,3 +1,4 @@
+use std::collections::VecDeque;
 use std::io::Write;
 
 use crossterm::style::Color;
@@ -8,12 +9,13 @@ use crossterm::{
 use eyre::{
     Result,
     WrapErr,
+    eyre,
 };
 use fig_os_shim::Context;
 use serde::Deserialize;
 
 use super::InvokeOutput;
-use crate::cli::chat::conversation_state::ConversationState;
+use crate::cli::chat::context::ContextManager;
 use crate::cli::issue::IssueCreator;
 
 #[derive(Debug, Clone, Deserialize)]
@@ -22,25 +24,36 @@ pub struct GhIssue {
     pub expected_behavior: Option<String>,
     pub actual_behavior: Option<String>,
     pub steps_to_reproduce: Option<String>,
+
+    #[serde(skip_deserializing)]
+    pub context: Option<GhIssueContext>,
 }
 
-pub struct GhIssueContext<'a> {
-    pub conversation_state: &'a ConversationState,
-    pub failed_request_ids: &'a Vec<String>,
+#[derive(Debug, Clone)]
+pub struct GhIssueContext {
+    pub context_manager: Option<ContextManager>,
+    pub transcript: VecDeque<String>,
+    pub failed_request_ids: Vec<String>,
 }
 
 /// Max amount of user chat + assistant recent chat messages to include in the issue.
 const MAX_TRANSCRIPT_LEN: usize = 10;
 
 impl GhIssue {
-    pub async fn invoke(&self, _updates: impl Write, context: GhIssueContext<'_>) -> Result<InvokeOutput> {
+    pub async fn invoke(&self, _updates: impl Write) -> Result<InvokeOutput> {
+        let Some(context) = self.context.as_ref() else {
+            return Err(eyre!(
+                "report_issue: Required tool context (GhIssueContext) not set by the program."
+            ));
+        };
+
         // Prepare additional details from the chat session
-        let additional_environment = [Self::get_request_ids(&context), Self::get_context(&context).await].join("\n\n");
+        let additional_environment = [Self::get_request_ids(context), Self::get_context(context).await].join("\n\n");
 
         // Add chat history to the actual behavior text.
         let actual_behavior = self.actual_behavior.as_ref().map_or_else(
-            || Self::get_transcript(&context),
-            |behavior| format!("{behavior}\n\n{}\n", Self::get_transcript(&context)),
+            || Self::get_transcript(context),
+            |behavior| format!("{behavior}\n\n{}\n", Self::get_transcript(context)),
         );
 
         let _ = IssueCreator {
@@ -57,9 +70,13 @@ impl GhIssue {
         Ok(Default::default())
     }
 
-    fn get_transcript(context: &GhIssueContext<'_>) -> String {
+    pub fn set_context(&mut self, context: GhIssueContext) {
+        self.context = Some(context);
+    }
+
+    fn get_transcript(context: &GhIssueContext) -> String {
         let mut transcript_str = String::from("```\n[chat-transcript]\n");
-        let transcript: Vec<String> = context.conversation_state.transcript
+        let transcript: Vec<String> = context.transcript
             .iter()
             .rev() // To take last N items
             .scan(0, |user_msg_count, line| {
@@ -89,7 +106,7 @@ impl GhIssue {
         transcript_str
     }
 
-    fn get_request_ids(context: &GhIssueContext<'_>) -> String {
+    fn get_request_ids(context: &GhIssueContext) -> String {
         format!(
             "[chat-failed_request_ids]\n{}",
             if context.failed_request_ids.is_empty() {
@@ -100,9 +117,9 @@ impl GhIssue {
         )
     }
 
-    async fn get_context(context: &GhIssueContext<'_>) -> String {
+    async fn get_context(context: &GhIssueContext) -> String {
         let mut ctx_str = "[chat-context]\n".to_string();
-        let Some(ctx_manager) = &context.conversation_state.context_manager else {
+        let Some(ctx_manager) = &context.context_manager else {
             ctx_str.push_str("No context available.");
             return ctx_str;
         };

crates/q_cli/src/cli/chat/tools/mod.rs

@@ -24,10 +24,7 @@ use fig_api_client::model::{
 use fig_os_shim::Context;
 use fs_read::FsRead;
 use fs_write::FsWrite;
-use gh_issue::{
-    GhIssue,
-    GhIssueContext,
-};
+use gh_issue::GhIssue;
 use serde::Deserialize;
 use use_aws::UseAws;
 
@@ -81,19 +78,13 @@ impl Tool {
     }
 
     /// Invokes the tool asynchronously
-    // TODO: Need to rework this to avoid passing in args meant for a single tool.
-    pub async fn invoke(
-        &self,
-        context: &Context,
-        updates: &mut impl Write,
-        gh_issue_context: GhIssueContext<'_>,
-    ) -> Result<InvokeOutput> {
+    pub async fn invoke(&self, context: &Context, updates: &mut impl Write) -> Result<InvokeOutput> {
         match self {
             Tool::FsRead(fs_read) => fs_read.invoke(context, updates).await,
             Tool::FsWrite(fs_write) => fs_write.invoke(context, updates).await,
             Tool::ExecuteBash(execute_bash) => execute_bash.invoke(updates).await,
             Tool::UseAws(use_aws) => use_aws.invoke(context, updates).await,
-            Tool::GhIssue(gh_issue) => gh_issue.invoke(updates, gh_issue_context).await,
+            Tool::GhIssue(gh_issue) => gh_issue.invoke(updates).await,
         }
     }