wafer-cli 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

wafer/wevin_cli.py

@@ -25,6 +25,7 @@ class StreamingChunkFrontend:
 
     Designed for programmatic consumption by extensions/UIs.
     Emits events in the format expected by wevin-extension handleWevinEvent:
+    - {type: 'session_start', session_id: '...', model: '...'}
     - {type: 'text_delta', delta: '...'}
     - {type: 'tool_call_start', tool_name: '...'}
     - {type: 'tool_call_end', tool_name: '...', args: {...}}
@@ -33,16 +34,31 @@ class StreamingChunkFrontend:
     - {type: 'error', error: '...'}
     """
 
-    def __init__(self) -> None:
+    def __init__(self, session_id: str | None = None, model: str | None = None) -> None:
         self._current_tool_call: dict | None = None
+        self._session_id = session_id
+        self._model = model
 
     def _emit(self, obj: dict) -> None:
         """Emit a single NDJSON line."""
         print(json.dumps(obj, ensure_ascii=False), flush=True)
 
     async def start(self) -> None:
-        """Initialize frontend."""
-        pass
+        """Initialize frontend and emit session_start if session_id is known."""
+        if self._session_id:
+            self._emit({
+                "type": "session_start",
+                "session_id": self._session_id,
+                "model": self._model,
+            })
+
+    def emit_session_start(self, session_id: str, model: str | None = None) -> None:
+        """Emit session_start event (for new sessions created during run)."""
+        self._emit({
+            "type": "session_start",
+            "session_id": session_id,
+            "model": model or self._model,
+        })
 
     async def stop(self) -> None:
         """Emit session_end event."""
@@ -335,31 +351,116 @@ def main(  # noqa: PLR0913, PLR0915
     single_turn: bool | None = None,  # None = use template default
     model: str | None = None,
     resume: str | None = None,
+    from_turn: int | None = None,
     tools: list[str] | None = None,
+    allow_spawn: bool = False,
+    max_tool_fails: int | None = None,
+    max_turns: int | None = None,
     template: str | None = None,
     template_args: dict[str, str] | None = None,
     corpus_path: str | None = None,
     list_sessions: bool = False,
+    get_session: str | None = None,
     json_output: bool = False,
-    # Legacy args (ignored)
-    problem: str | None = None,
-    reference: str | None = None,
-    **kwargs: object,
 ) -> None:
     """Run wevin agent in-process via rollouts."""
+    import trio
+    from dataclasses import asdict
+
     from wafer_core.rollouts import FileSessionStore
 
+    session_store = FileSessionStore()
+
+    # Handle --get-session: load session by ID and print
+    if get_session:
+        async def _get_session() -> None:
+            try:
+                session, err = await session_store.get(get_session)
+                if err or not session:
+                    if json_output:
+                        print(json.dumps({"error": err or f"Session {get_session} not found"}))
+                        sys.exit(1)
+                    else:
+                        print(f"Error: {err or 'Session not found'}", file=sys.stderr)
+                        sys.exit(1)
+
+                if json_output:
+                    # Serialize messages to dicts
+                    try:
+                        messages_data = [asdict(msg) for msg in session.messages]
+                    except Exception as e:
+                        # If serialization fails, return error
+                        error_msg = f"Failed to serialize messages: {e}"
+                        print(json.dumps({"error": error_msg}))
+                        sys.exit(1)
+                    
+                    print(json.dumps({
+                        "session_id": session.session_id,
+                        "status": session.status.value,
+                        "model": session.endpoint.model if session.endpoint else None,
+                        "created_at": session.created_at,
+                        "updated_at": session.updated_at,
+                        "messages": messages_data,
+                        "tags": session.tags,
+                    }))
+                else:
+                    print(f"Session: {session.session_id}")
+                    print(f"Status: {session.status.value}")
+                    print(f"Messages: {len(session.messages)}")
+                    for i, msg in enumerate(session.messages):
+                        # Fail fast if message can't be converted to string - corrupted data is a bug
+                        content_preview = str(msg.content)[:100] if msg.content else ""
+                        print(f"  [{i}] {msg.role}: {content_preview}...")
+            except KeyboardInterrupt:
+                # User cancelled - exit cleanly
+                sys.exit(130)  # Standard exit code for SIGINT
+            except Exception as e:
+                # Any other error - log and exit with error
+                error_msg = f"Failed to load session {get_session}: {e}"
+                if json_output:
+                    print(json.dumps({"error": error_msg}))
+                else:
+                    print(f"Error: {error_msg}", file=sys.stderr)
+                sys.exit(1)
+
+        try:
+            trio.run(_get_session)
+        except KeyboardInterrupt:
+            sys.exit(130)
+        except Exception as e:
+            error_msg = f"Failed to run session loader: {e}"
+            if json_output:
+                print(json.dumps({"error": error_msg}))
+            else:
+                print(f"Error: {error_msg}", file=sys.stderr)
+            sys.exit(1)
+        return
+
     # Handle --list-sessions: show recent sessions and exit
     if list_sessions:
-        session_store = FileSessionStore()
-        sessions = session_store.list_sync(limit=20)
-        if not sessions:
-            print("No sessions found.")
-        else:
-            print("Recent sessions:")
+        sessions = session_store.list_sync(limit=50)
+        if json_output:
+            # Return metadata only - messages loaded on-demand via --get-session
+            sessions_data = []
             for s in sessions:
-                preview = _get_session_preview(s)
-                print(f"  {s.session_id}  {preview}")
+                sessions_data.append({
+                    "session_id": s.session_id,
+                    "status": s.status.value,
+                    "model": s.endpoint.model if s.endpoint else None,
+                    "created_at": s.created_at if hasattr(s, "created_at") else None,
+                    "updated_at": s.updated_at if hasattr(s, "updated_at") else None,
+                    "message_count": len(s.messages),
+                    "preview": _get_session_preview(s),
+                })
+            print(json.dumps({"sessions": sessions_data}))
+        else:
+            if not sessions:
+                print("No sessions found.")
+            else:
+                print("Recent sessions:")
+                for s in sessions:
+                    preview = _get_session_preview(s)
+                    print(f"  {s.session_id}  {preview}")
         return
 
     # Emit early event for JSON mode before heavy imports
@@ -367,8 +468,7 @@ def main(  # noqa: PLR0913, PLR0915
     if json_output:
         print(json.dumps({"type": "initializing"}), flush=True)
 
-    import trio
-    from wafer_core.rollouts import FileSessionStore, Message, Trajectory
+    from wafer_core.rollouts import Message, Trajectory
     from wafer_core.rollouts.frontends import NoneFrontend, RunnerConfig, run_interactive
 
     _setup_logging()
@@ -444,7 +544,9 @@ def main(  # noqa: PLR0913, PLR0915
                 )
             else:
                 if json_output:
-                    frontend = StreamingChunkFrontend()
+                    # Emit session_start if we have a session_id (from --resume)
+                    model_name = endpoint.model if hasattr(endpoint, 'model') else None
+                    frontend = StreamingChunkFrontend(session_id=session_id, model=model_name)
                 else:
                     frontend = NoneFrontend(show_tool_calls=True, show_thinking=False)
                 config = RunnerConfig(
@@ -455,6 +557,13 @@ def main(  # noqa: PLR0913, PLR0915
                     hide_session_info=True,  # We print our own resume command
                 )
                 states = await run_interactive(trajectory, endpoint, frontend, environment, config)
+                # Emit session_start for new sessions (if session_id was None and we got one)
+                # Check first state to emit as early as possible
+                if json_output and isinstance(frontend, StreamingChunkFrontend):
+                    first_session_id = states[0].session_id if states and states[0].session_id else None
+                    if first_session_id and not session_id:  # New session created
+                        model_name = endpoint.model if hasattr(endpoint, 'model') else None
+                        frontend.emit_session_start(first_session_id, model_name)
                 # Print resume command with full wafer agent prefix
                 if states and states[-1].session_id:
                     print(f"\nResume with: wafer agent --resume {states[-1].session_id}")