From 64bb9e9ca6df5873a939d36f2e95e375cb69121b Mon Sep 17 00:00:00 2001
From: Salman Mohammed <smohammed@squareup.com>
Date: Wed, 5 Feb 2025 10:53:35 -0500
Subject: [PATCH] fix: add logs for goose run (#1085)

---
 crates/goose-cli/src/main.rs    | 2 +-
 crates/goose-cli/src/session.rs | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/crates/goose-cli/src/main.rs b/crates/goose-cli/src/main.rs
index 84f0a9cac..d7b154d4f 100644
--- a/crates/goose-cli/src/main.rs
+++ b/crates/goose-cli/src/main.rs
@@ -180,7 +180,6 @@ async fn main() -> Result<()> {
         }) => {
             let mut session = build_session(name, resume, extension, builtin).await;
             setup_logging(session.session_file().file_stem().and_then(|s| s.to_str()))?;
-
             let _ = session.start().await;
             return Ok(());
         }
@@ -211,6 +210,7 @@ async fn main() -> Result<()> {
                 stdin
             };
             let mut session = build_session(name, resume, extension, builtin).await;
+            setup_logging(session.session_file().file_stem().and_then(|s| s.to_str()))?;
             let _ = session.headless_start(contents.clone()).await;
             return Ok(());
         }
diff --git a/crates/goose-cli/src/session.rs b/crates/goose-cli/src/session.rs
index 7b7d9644e..dfd0b0810 100644
--- a/crates/goose-cli/src/session.rs
+++ b/crates/goose-cli/src/session.rs
@@ -309,6 +309,9 @@ We've removed the conversation up to the most recent user message
     }
 
     async fn close_session(&mut self) {
+        let usage = self.agent.usage().await;
+        log_usage(self.session_file.to_string_lossy().to_string(), usage);
+
         self.prompt.render(raw_message(
             format!(
                 "Closing session. Recorded to {}\n",
@@ -317,8 +320,6 @@ We've removed the conversation up to the most recent user message
             .as_str(),
         ));
         self.prompt.close();
-        let usage = self.agent.usage().await;
-        log_usage(self.session_file.to_string_lossy().to_string(), usage);
     }
 
     pub fn session_file(&self) -> PathBuf {