aimo_note_agent/
lib.rs

1use std::sync::Arc;
2
3use amico_core::{
4    Agent,
5    types::{Chat, ChatMessage},
6};
7use tokio::sync::Mutex;
8use tokio_with_wasm::alias as tokio;
9use wasm_bindgen::prelude::*;
10use wasm_bindgen_futures::spawn_local;
11
12mod agent;
13mod log;
14mod service;
15mod note;
16
17use agent::{AppStrategy, ChatHandler, create_agent};
18use note::Note;
19
20use crate::agent::ChatContext;
21
22/// A WASM-bindgen compatible message structure that can be converted to ChatMessage.
23#[wasm_bindgen]
24#[derive(Clone, Debug)]
25pub struct Message {
26    content: String,
27    role: String,
28}
29
30#[wasm_bindgen]
31impl Message {
32    #[wasm_bindgen(constructor)]
33    pub fn new(content: String, role: String) -> Message {
34        Message { content, role }
35    }
36
37    #[wasm_bindgen(getter)]
38    pub fn content(&self) -> String {
39        self.content.clone()
40    }
41
42    #[wasm_bindgen(getter)]
43    pub fn role(&self) -> String {
44        self.role.clone()
45    }
46
47    #[wasm_bindgen(setter)]
48    pub fn set_content(&mut self, content: String) {
49        self.content = content;
50    }
51
52    #[wasm_bindgen(setter)]
53    pub fn set_role(&mut self, role: String) {
54        self.role = role;
55    }
56}
57
58impl From<Message> for ChatMessage {
59    fn from(message: Message) -> Self {
60        ChatMessage {
61            content: message.content,
62            role: message.role,
63        }
64    }
65}
66
67/// The WASM runtime for the agent.
68#[wasm_bindgen]
69pub struct AgentWasmRuntime {
70    agent: Option<Agent<AppStrategy>>,
71    chat_handler: Arc<Mutex<ChatHandler>>,
72    running: bool,
73}
74
75#[wasm_bindgen]
76impl AgentWasmRuntime {
77    #[wasm_bindgen(constructor)]
78    pub fn new(jwt: String) -> AgentWasmRuntime {
79        let (agent, chat_handler) = create_agent(jwt);
80
81        AgentWasmRuntime {
82            agent: Some(agent),
83            chat_handler: Arc::new(Mutex::new(chat_handler)),
84            running: false,
85        }
86    }
87
88    #[wasm_bindgen]
89    pub fn start(&mut self) {
90        if self.running {
91            tracing::warn!("Agent is already running");
92            return;
93        }
94
95        if let Some(mut agent) = self.agent.take() {
96            let _chat_handler = self.chat_handler.clone();
97            spawn_local(async move {
98                tracing::info!("Starting agent runtime");
99                agent.run().await;
100            });
101            self.running = true;
102            tracing::info!("Agent runtime started");
103        }
104    }
105
106    #[wasm_bindgen]
107    pub async fn chat(&self, messages: Vec<Message>, cursor_position: usize, note: JsValue) -> Result<JsValue, JsValue> {
108        if !self.running {
109            return Err(JsValue::from_str(
110                "Agent is not running. Call start() first.",
111            ));
112        }
113
114        // Parse the note from the JS value.
115        let note: Note = serde_wasm_bindgen::from_value(note)?;
116
117        // Convert Vec<Message> to Vec<ChatMessage>
118        let chat_messages: Vec<ChatMessage> = messages.into_iter().map(|msg| msg.into()).collect();
119
120        let chat = Chat {
121            messages: chat_messages,
122
123            // We don't use session_id here
124            session_id: 0,
125        };
126
127        let mut handler = self.chat_handler.lock().await;
128        match handler.chat(chat, &ChatContext {
129            note,
130            cursor_position,
131        }).await {
132            Ok(action) => Ok(serde_wasm_bindgen::to_value(&action)?),
133            Err(e) => Err(JsValue::from_str(&format!("Chat error: {}", e))),
134        }
135    }
136
137    #[wasm_bindgen]
138    pub fn is_running(&self) -> bool {
139        self.running
140    }
141}
142
143/// Initialize the WASM module.
144#[wasm_bindgen(start)]
145pub fn start() {
146    log::init();
147    tracing::info!("WASM module initialized");
148}
aimo_note_agent/lib.rs

aimo_note_agent/
lib.rs
