Skip to main content

coderlib/integration/
ai_assistant.rs

1//! AI Assistant Interface for Edit integration
2//!
3//! This module provides the user interface components and interaction
4//! logic for AI assistance within Microsoft Edit.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::{mpsc, RwLock};
10use tracing::{debug, info};
11
12use crate::core::{CoderLib, CoderLibError, CodeRequest, CodeResponse};
13use crate::integration::{
14    EditHost, ContextGatherer, ContextConfig, GatheredContext,
15    HostCommand
16};
17use crate::storage::{Message, MessageRole, MessageContent};
18
19/// AI Assistant interface for Edit
20pub struct AIAssistant {
21    /// CoderLib instance
22    coderlib: Arc<CoderLib>,
23    /// Edit host integration
24    edit_host: Arc<EditHost>,
25    /// Context gatherer
26    context_gatherer: ContextGatherer,
27    /// Active sessions
28    sessions: Arc<RwLock<HashMap<String, AssistantSession>>>,
29    /// UI state
30    ui_state: Arc<RwLock<UIState>>,
31    /// Command sender to Edit
32    command_sender: Option<mpsc::UnboundedSender<HostCommand>>,
33}
34
35/// AI Assistant session
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct AssistantSession {
38    /// Session ID
39    pub id: String,
40    /// Session type
41    pub session_type: SessionType,
42    /// Conversation history
43    pub messages: Vec<Message>,
44    /// Current context
45    pub context: Option<GatheredContext>,
46    /// Session state
47    pub state: SessionState,
48    /// Created timestamp
49    pub created_at: std::time::SystemTime,
50    /// Last activity timestamp
51    pub last_activity: std::time::SystemTime,
52}
53
54/// Types of AI assistant sessions
55#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
56pub enum SessionType {
57    /// General code assistance
58    CodeAssistance,
59    /// Code explanation
60    CodeExplanation,
61    /// Code refactoring
62    CodeRefactoring,
63    /// Bug fixing
64    BugFix,
65    /// Code review
66    CodeReview,
67    /// Documentation generation
68    Documentation,
69    /// Test generation
70    TestGeneration,
71    /// Performance optimization
72    Performance,
73}
74
75/// Session state
76#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
77pub enum SessionState {
78    /// Session is active and ready for input
79    Active,
80    /// AI is processing a request
81    Processing,
82    /// Waiting for user response
83    WaitingForUser,
84    /// Session is paused
85    Paused,
86    /// Session completed successfully
87    Completed,
88    /// Session ended with error
89    Error(String),
90}
91
92/// UI state for the AI assistant
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct UIState {
95    /// Whether the assistant panel is visible
96    pub panel_visible: bool,
97    /// Current active session ID
98    pub active_session: Option<String>,
99    /// Panel size and position
100    pub panel_bounds: PanelBounds,
101    /// UI preferences
102    pub preferences: UIPreferences,
103}
104
105/// Panel bounds and positioning
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct PanelBounds {
108    /// Panel width (percentage of editor width)
109    pub width_percent: f32,
110    /// Panel height (percentage of editor height)
111    pub height_percent: f32,
112    /// Panel position
113    pub position: PanelPosition,
114}
115
116/// Panel position options
117#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
118pub enum PanelPosition {
119    /// Right side of editor
120    Right,
121    /// Left side of editor
122    Left,
123    /// Bottom of editor
124    Bottom,
125    /// Floating window
126    Floating { x: i32, y: i32 },
127}
128
129/// UI preferences
130#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct UIPreferences {
132    /// Theme (light/dark)
133    pub theme: String,
134    /// Font size
135    pub font_size: u32,
136    /// Show token usage
137    pub show_token_usage: bool,
138    /// Auto-apply simple changes
139    pub auto_apply_simple: bool,
140    /// Show context preview
141    pub show_context_preview: bool,
142    /// Maximum visible messages
143    pub max_visible_messages: usize,
144}
145
146/// AI Assistant request
147#[derive(Debug, Clone, Serialize, Deserialize)]
148pub struct AssistantRequest {
149    /// Session ID (optional, creates new session if not provided)
150    pub session_id: Option<String>,
151    /// Session type
152    pub session_type: SessionType,
153    /// User message
154    pub message: String,
155    /// Whether to include context
156    pub include_context: bool,
157    /// Additional context hints
158    pub context_hints: Vec<String>,
159}
160
161/// AI Assistant response
162#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct AssistantResponse {
164    /// Session ID
165    pub session_id: String,
166    /// Response message
167    pub message: String,
168    /// Suggested actions
169    pub suggested_actions: Vec<SuggestedAction>,
170    /// Token usage information
171    pub token_usage: Option<TokenUsage>,
172    /// Response metadata
173    pub metadata: HashMap<String, serde_json::Value>,
174}
175
176/// Suggested action from AI
177#[derive(Debug, Clone, Serialize, Deserialize)]
178pub struct SuggestedAction {
179    /// Action type
180    pub action_type: ActionType,
181    /// Action description
182    pub description: String,
183    /// Action data
184    pub data: serde_json::Value,
185    /// Confidence score (0.0 to 1.0)
186    pub confidence: f64,
187}
188
189/// Types of suggested actions
190#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
191pub enum ActionType {
192    /// Apply code changes
193    ApplyChanges,
194    /// Open file
195    OpenFile,
196    /// Create new file
197    CreateFile,
198    /// Run command
199    RunCommand,
200    /// Show documentation
201    ShowDocumentation,
202    /// Navigate to definition
203    NavigateToDefinition,
204    /// Add import/dependency
205    AddImport,
206    /// Generate tests
207    GenerateTests,
208}
209
210/// Token usage information
211#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct TokenUsage {
213    /// Input tokens
214    pub input_tokens: u32,
215    /// Output tokens
216    pub output_tokens: u32,
217    /// Total tokens
218    pub total_tokens: u32,
219    /// Estimated cost
220    pub estimated_cost: Option<f64>,
221}
222
223impl Default for UIState {
224    fn default() -> Self {
225        Self {
226            panel_visible: false,
227            active_session: None,
228            panel_bounds: PanelBounds {
229                width_percent: 40.0,
230                height_percent: 60.0,
231                position: PanelPosition::Right,
232            },
233            preferences: UIPreferences {
234                theme: "dark".to_string(),
235                font_size: 14,
236                show_token_usage: true,
237                auto_apply_simple: false,
238                show_context_preview: true,
239                max_visible_messages: 50,
240            },
241        }
242    }
243}
244
245impl AIAssistant {
246    /// Create a new AI assistant
247    pub fn new(
248        coderlib: Arc<CoderLib>,
249        edit_host: Arc<EditHost>,
250        context_config: ContextConfig,
251    ) -> Self {
252        let context_gatherer = ContextGatherer::new(context_config);
253        
254        Self {
255            coderlib,
256            edit_host,
257            context_gatherer,
258            sessions: Arc::new(RwLock::new(HashMap::new())),
259            ui_state: Arc::new(RwLock::new(UIState::default())),
260            command_sender: None,
261        }
262    }
263
264    /// Set the command sender for communicating with Edit
265    pub fn set_command_sender(&mut self, sender: mpsc::UnboundedSender<HostCommand>) {
266        self.command_sender = Some(sender);
267    }
268
269    /// Show the AI assistant panel
270    pub async fn show_panel(&self) -> Result<(), CoderLibError> {
271        let mut ui_state = self.ui_state.write().await;
272        ui_state.panel_visible = true;
273        
274        self.send_command(HostCommand::ShowDialog {
275            title: "AI Assistant".to_string(),
276            message: "AI Assistant is ready to help!".to_string(),
277            buttons: vec!["Ask Question".to_string(), "Explain Code".to_string(), "Refactor".to_string(), "Close".to_string()],
278        }).await?;
279        
280        info!("AI Assistant panel shown");
281        Ok(())
282    }
283
284    /// Hide the AI assistant panel
285    pub async fn hide_panel(&self) -> Result<(), CoderLibError> {
286        let mut ui_state = self.ui_state.write().await;
287        ui_state.panel_visible = false;
288        
289        info!("AI Assistant panel hidden");
290        Ok(())
291    }
292
293    /// Process an AI assistant request
294    pub async fn process_request(&self, request: AssistantRequest) -> Result<AssistantResponse, CoderLibError> {
295        info!("Processing AI assistant request: {:?}", request.session_type);
296
297        // Get or create session
298        let session_id = if let Some(ref id) = request.session_id {
299            id.clone()
300        } else {
301            self.create_new_session(request.session_type.clone()).await?
302        };
303
304        // Update session state to processing
305        self.update_session_state(&session_id, SessionState::Processing).await?;
306
307        // Gather context if requested
308        let context = if request.include_context {
309            let state = self.edit_host.get_state();
310            Some(self.context_gatherer.gather_context(&state, self.edit_host.as_ref()).await?)
311        } else {
312            None
313        };
314
315        // Build the code request
316        let code_request = self.build_code_request(&request, &context).await?;
317
318        // Process with CoderLib
319        let mut code_response_receiver = self.coderlib.process_request(code_request).await?;
320
321        // Get the first response (simplified for now)
322        let code_response = if let Ok(response) = code_response_receiver.recv().await {
323            response
324        } else {
325            return Err(CoderLibError::Integration(
326                crate::core::error::IntegrationError::OperationFailed("No response received".to_string())
327            ));
328        };
329
330        // Convert to assistant response
331        let assistant_response = self.build_assistant_response(&session_id, &code_response, &context).await?;
332
333        // Update session with new messages
334        self.update_session_messages(&session_id, &request.message, &assistant_response.message).await?;
335
336        // Update session state
337        self.update_session_state(&session_id, SessionState::WaitingForUser).await?;
338
339        info!("AI assistant request processed successfully");
340        Ok(assistant_response)
341    }
342
343    /// Create a new assistant session
344    async fn create_new_session(&self, session_type: SessionType) -> Result<String, CoderLibError> {
345        let session_id = uuid::Uuid::new_v4().to_string();
346        let now = std::time::SystemTime::now();
347        
348        let session = AssistantSession {
349            id: session_id.clone(),
350            session_type,
351            messages: Vec::new(),
352            context: None,
353            state: SessionState::Active,
354            created_at: now,
355            last_activity: now,
356        };
357
358        let mut sessions = self.sessions.write().await;
359        sessions.insert(session_id.clone(), session);
360
361        // Update UI state
362        let mut ui_state = self.ui_state.write().await;
363        ui_state.active_session = Some(session_id.clone());
364
365        debug!("Created new AI assistant session: {}", session_id);
366        Ok(session_id)
367    }
368
369    /// Update session state
370    async fn update_session_state(&self, session_id: &str, new_state: SessionState) -> Result<(), CoderLibError> {
371        let mut sessions = self.sessions.write().await;
372        if let Some(session) = sessions.get_mut(session_id) {
373            session.state = new_state;
374            session.last_activity = std::time::SystemTime::now();
375        }
376        Ok(())
377    }
378
379    /// Update session messages
380    async fn update_session_messages(&self, session_id: &str, user_message: &str, assistant_message: &str) -> Result<(), CoderLibError> {
381        let mut sessions = self.sessions.write().await;
382        if let Some(session) = sessions.get_mut(session_id) {
383            session.messages.push(Message {
384                id: uuid::Uuid::new_v4().to_string(),
385                session_id: session_id.to_string(),
386                role: MessageRole::User,
387                content: MessageContent::Text(user_message.to_string()),
388                timestamp: chrono::Utc::now(),
389                metadata: serde_json::Value::Null,
390            });
391            session.messages.push(Message {
392                id: uuid::Uuid::new_v4().to_string(),
393                session_id: session_id.to_string(),
394                role: MessageRole::Assistant,
395                content: MessageContent::Text(assistant_message.to_string()),
396                timestamp: chrono::Utc::now(),
397                metadata: serde_json::Value::Null,
398            });
399            session.last_activity = std::time::SystemTime::now();
400        }
401        Ok(())
402    }
403
404    /// Build code request from assistant request
405    async fn build_code_request(&self, request: &AssistantRequest, context: &Option<GatheredContext>) -> Result<CodeRequest, CoderLibError> {
406        let mut prompt = String::new();
407
408        // Add session type specific prompt
409        match request.session_type {
410            SessionType::CodeExplanation => {
411                prompt.push_str("Please explain the following code:\n\n");
412            }
413            SessionType::CodeRefactoring => {
414                prompt.push_str("Please refactor the following code to improve its quality, readability, and performance:\n\n");
415            }
416            SessionType::BugFix => {
417                prompt.push_str("Please help me identify and fix bugs in the following code:\n\n");
418            }
419            SessionType::CodeReview => {
420                prompt.push_str("Please review the following code and provide feedback:\n\n");
421            }
422            SessionType::Documentation => {
423                prompt.push_str("Please generate documentation for the following code:\n\n");
424            }
425            SessionType::TestGeneration => {
426                prompt.push_str("Please generate comprehensive tests for the following code:\n\n");
427            }
428            SessionType::Performance => {
429                prompt.push_str("Please analyze and optimize the performance of the following code:\n\n");
430            }
431            SessionType::CodeAssistance => {
432                prompt.push_str("Please help me with the following code question:\n\n");
433            }
434        }
435
436        // Add user message
437        prompt.push_str(&request.message);
438        prompt.push_str("\n\n");
439
440        // Add context if available
441        if let Some(ctx) = context {
442            prompt.push_str("## Context\n");
443            prompt.push_str(&self.context_gatherer.format_context_for_ai(ctx));
444        }
445
446        Ok(CodeRequest {
447            content: prompt,
448            attachments: Vec::new(),
449            model: None,
450            context: crate::core::RequestContext {
451                current_file: None,
452                cursor_position: None,
453                selection: None,
454                project_root: None,
455                open_files: Vec::new(),
456            },
457            session_id: uuid::Uuid::new_v4().to_string(),
458        })
459    }
460
461    /// Build assistant response from code response
462    async fn build_assistant_response(&self, session_id: &str, code_response: &CodeResponse, _context: &Option<GatheredContext>) -> Result<AssistantResponse, CoderLibError> {
463        let suggested_actions = self.extract_suggested_actions(&code_response.content).await;
464
465        Ok(AssistantResponse {
466            session_id: session_id.to_string(),
467            message: code_response.content.clone(),
468            suggested_actions,
469            token_usage: None, // TODO: Extract from code_response when available
470            metadata: HashMap::new(),
471        })
472    }
473
474    /// Extract suggested actions from AI response
475    async fn extract_suggested_actions(&self, _response_content: &str) -> Vec<SuggestedAction> {
476        // This is a simplified implementation
477        // In a real implementation, we would parse the AI response for actionable items
478        vec![]
479    }
480
481    /// Send a command to Edit
482    async fn send_command(&self, command: HostCommand) -> Result<(), CoderLibError> {
483        if let Some(sender) = &self.command_sender {
484            sender.send(command)
485                .map_err(|e| CoderLibError::Integration(
486                    crate::core::error::IntegrationError::OperationFailed(
487                        format!("Failed to send command: {}", e)
488                    )
489                ))?;
490        }
491        Ok(())
492    }
493
494    /// Get current UI state
495    pub async fn get_ui_state(&self) -> UIState {
496        self.ui_state.read().await.clone()
497    }
498
499    /// Update UI preferences
500    pub async fn update_preferences(&self, preferences: UIPreferences) -> Result<(), CoderLibError> {
501        let mut ui_state = self.ui_state.write().await;
502        ui_state.preferences = preferences;
503        Ok(())
504    }
505
506    /// Get session by ID
507    pub async fn get_session(&self, session_id: &str) -> Option<AssistantSession> {
508        let sessions = self.sessions.read().await;
509        sessions.get(session_id).cloned()
510    }
511
512    /// List all sessions
513    pub async fn list_sessions(&self) -> Vec<AssistantSession> {
514        let sessions = self.sessions.read().await;
515        sessions.values().cloned().collect()
516    }
517
518    /// Close session
519    pub async fn close_session(&self, session_id: &str) -> Result<(), CoderLibError> {
520        let mut sessions = self.sessions.write().await;
521        sessions.remove(session_id);
522        
523        // Update UI state if this was the active session
524        let mut ui_state = self.ui_state.write().await;
525        if ui_state.active_session.as_ref() == Some(&session_id.to_string()) {
526            ui_state.active_session = None;
527        }
528        
529        Ok(())
530    }
531}