(chore): add tracing option to run and group traces under session name

src/goose/cli/main.py

@@ -9,7 +9,7 @@
 
 from goose.cli.config import SESSIONS_PATH
 from goose.cli.session import Session
-from goose.toolkit.utils import render_template, parse_plan
+from goose.toolkit.utils import parse_plan, render_template
 from goose.utils import load_plugins
 from goose.utils.autocomplete import SUPPORTED_SHELLS, setup_autocomplete
 from goose.utils.session_file import list_sorted_session_files
@@ -208,7 +208,10 @@ def session_resume(name: Optional[str], profile: str, log_level: str) -> None:
 @click.option("--profile")
 @click.option("--log-level", type=LOG_CHOICE, default="INFO")
 @click.option("--resume-session", is_flag=True, help="Resume the last session if available")
-def run(message_file: Optional[str], profile: str, log_level: str, resume_session: bool = False) -> None:
+@click.option("--tracing", is_flag=True, required=False)
+def run(
+    message_file: Optional[str], profile: str, log_level: str, resume_session: bool = False, tracing: bool = False
+) -> None:
     """Run a single-pass session with a message from a markdown input file"""
     if message_file:
         with open(message_file, "r") as f:
@@ -220,9 +223,9 @@ def run(message_file: Optional[str], profile: str, log_level: str, resume_sessio
         session_files = get_session_files()
         if session_files:
             name = list(session_files.keys())[0]
-            session = Session(name=name, profile=profile, log_level=log_level)
+            session = Session(name=name, profile=profile, log_level=log_level, tracing=tracing)
     else:
-        session = Session(profile=profile, log_level=log_level)
+        session = Session(profile=profile, log_level=log_level, tracing=tracing)
     session.single_pass(initial_message=initial_message)
 
 

src/goose/cli/session.py

@@ -3,9 +3,9 @@
 from pathlib import Path
 from typing import Optional
 
-from langfuse.decorators import langfuse_context
 from exchange import Message, Text, ToolResult, ToolUse
-from exchange.langfuse_wrapper import observe_wrapper, auth_check
+from exchange.langfuse_wrapper import auth_check, observe_wrapper
+from langfuse.decorators import langfuse_context
 from rich import print
 from rich.markdown import Markdown
 from rich.panel import Panel
@@ -21,7 +21,7 @@
 from goose.utils import droid, load_plugins
 from goose.utils._cost_calculator import get_total_cost_message
 from goose.utils._create_exchange import create_exchange
-from goose.utils.session_file import is_empty_session, is_existing_session, read_or_create_file, log_messages
+from goose.utils.session_file import is_empty_session, is_existing_session, log_messages, read_or_create_file
 
 RESUME_MESSAGE = "I see we were interrupted. How can I help you?"
 
@@ -150,7 +150,6 @@ def single_pass(self, initial_message: str) -> None:
         profile = self.profile_name or "default"
         print(f"[dim]starting session | name:[cyan]{self.name}[/]  profile:[cyan]{profile}[/]")
         print(f"[dim]saving to {self.session_file_path}")
-        print()
 
         # Process initial message
         message = Message.user(initial_message)
@@ -202,6 +201,9 @@ def run(self, new_session: bool = True) -> None:
     @observe_wrapper()
     def reply(self) -> None:
         """Reply to the last user message, calling tools as needed"""
+        # group all traces under the same session
+        langfuse_context.update_current_trace(session_id=self.name)
+
         # These are the *raw* messages, before the moderator rewrites things
         committed = [self.exchange.messages[-1]]