diff --git a/src/agent.rs b/src/agent.rs
index 2bbc042..0aebd4e 100644
--- a/src/agent.rs
+++ b/src/agent.rs
@@ -1,4 +1,5 @@
 use chrono::Utc;
+use std::time::{Duration, Instant};
 
 use crate::api::{ChatRequest, ChatResponse, Message, Tool};
 use crate::tools;
@@ -68,9 +69,22 @@ impl Agent {
     }
 
     pub async fn run(&mut self) -> Result<(String, Option<String>), Box<dyn std::error::Error>> {
-        let mut file_written = false;
+        let mut finished = false;
+        let start_time = Instant::now();
+        let max_duration = Duration::from_secs(120);
+        let max_turns = 20;
+        let mut turns = 0;
 
-        while !file_written {
+        while !finished {
+            if start_time.elapsed() > max_duration {
+                return Err("Agent run timed out".into());
+            }
+
+            if turns >= max_turns {
+                return Err("Agent run exceeded max turns".into());
+            }
+
+            turns += 1;
             let request = ChatRequest {
                 model: "kimi-k2.5".to_string(),
                 messages: self.messages.clone(),
@@ -100,7 +114,12 @@ impl Agent {
             }
 
             let chat_response: ChatResponse = response.json().await?;
-            let assistant_message = chat_response.choices.get(0).unwrap().message.clone();
+            let assistant_message = chat_response
+                .choices
+                .get(0)
+                .ok_or("Missing assistant response")?
+                .message
+                .clone();
 
             self.messages.push(assistant_message.clone());
 
@@ -128,16 +147,16 @@ impl Agent {
 
                     self.messages.push(tool_message);
                     if written {
-                        file_written = true;
+                        finished = true;
                     }
                 }
                 // Continue the loop to send tool results back
                 continue;
             }
 
-            // No more tool calls from assistant, but we only exit if file was written
-            if !file_written {
-                self.log("--- Assistant didn't use write_file yet. Waiting for next turn... ---");
+            // No tool calls from assistant, but we only exit if the task was finished
+            if !finished {
+                self.log("--- Assistant didn't finish yet. Waiting for next turn... ---");
             }
         }