Merge pull request #120 from tower/feature/tow-771-cli-logs-are-truncated

TOW-771 - logs being truncated

Author: socksy

crates/tower-cmd/src/apps.rs

@@ -3,7 +3,7 @@ use config::Config;
 
 use tower_api::models::Run;
 
-use crate::{api, output, util::dates};
+use crate::{api, output};
 
 pub fn apps_cmd() -> Command {
     Command::new("apps")
@@ -51,8 +51,7 @@ pub async fn do_logs(config: Config, cmd: &ArgMatches) {
 
     if let Ok(resp) = api::describe_run_logs(&config, &name, seq).await {
         for line in resp.log_lines {
-            let ts = dates::format_str(&line.reported_at);
-            output::log_line(&ts, &line.content, output::LogLineType::Remote);
+            output::remote_log_event(&line);
         }
     }
 }

crates/tower-cmd/src/mcp.rs

@@ -25,6 +25,7 @@ use rmcp::{
 use rsa::pkcs1::DecodeRsaPublicKey;
 use serde::Deserialize;
 use serde_json::{json, Value};
+use tokio::net::TcpListener;
 use tower_api::apis::Error as ApiError;
 
 use crate::{api, deploy, run, towerfile_gen::TowerfileGenerator, Config, Error};
@@ -210,7 +211,7 @@ async fn run_http_server(config: Config, port: u16) -> Result<(), Error> {
     );
 
     let router = Router::new().nest_service("/mcp", service);
-    let listener = tokio::net::TcpListener::bind(&bind_addr).await?;
+    let listener = TcpListener::bind(&bind_addr).await?;
     axum::serve(listener, router)
         .with_graceful_shutdown(async {
             tokio::signal::ctrl_c().await.unwrap();

crates/tower-cmd/src/output.rs

@@ -127,6 +127,11 @@ pub fn log_line(timestamp: &str, message: &str, t: LogLineType) {
     write(&line);
 }
 
+pub fn remote_log_event(log: &tower_api::models::RunLogLine) {
+    let ts = crate::util::dates::format_str(&log.reported_at);
+    log_line(&ts, &log.content, LogLineType::Remote);
+}
+
 pub fn package_error(err: tower_package::Error) {
     let msg = match err {
         tower_package::Error::NoManifest => "No manifest was found".to_string(),

crates/tower-cmd/src/run.rs

@@ -8,7 +8,11 @@ use tower_package::{Package, PackageSpec};
 use tower_runtime::{local::LocalApp, App, AppLauncher, OutputReceiver, Status};
 use tower_telemetry::{debug, Context};
 
-use tokio::sync::{mpsc::unbounded_channel, oneshot};
+use tokio::sync::{
+    mpsc::{unbounded_channel, Receiver as MpscReceiver},
+    oneshot::{self, Receiver as OneshotReceiver},
+};
+use tokio::time::{sleep, timeout, Duration};
 
 use crate::{api, output, util::dates};
 
@@ -243,8 +247,50 @@ pub async fn do_run_remote(
     }
 }
 
+async fn stream_logs_until_complete(
+    mut log_stream: MpscReceiver<api::LogStreamEvent>,
+    mut run_complete: OneshotReceiver<Run>,
+    enable_ctrl_c: bool,
+    run_link: &str,
+) -> Result<Option<Run>, Error> {
+    loop {
+        tokio::select! {
+            event = log_stream.recv() => match event {
+                Some(api::LogStreamEvent::EventLog(log)) => {
+                    output::remote_log_event(&log);
+                },
+                None => return Ok(None),
+                _ => {},
+            },
+            res = &mut run_complete => {
+                let completed_run = res?;
+                drain_remaining_logs(log_stream).await;
+                return Ok(Some(completed_run));
+            },
+            _ = tokio::signal::ctrl_c(), if enable_ctrl_c => {
+                output::write("Received Ctrl+C, stopping log streaming...\n");
+                output::write("Note: The run will continue in Tower cloud\n");
+                output::write(&format!("  See more: {}\n", run_link));
+                return Ok(None);
+            },
+        }
+    }
+}
+
+async fn drain_remaining_logs(mut log_stream: MpscReceiver<api::LogStreamEvent>) {
+    let drain_duration = Duration::from_secs(5);
+    let _ = timeout(drain_duration, async {
+        while let Some(event) = log_stream.recv().await {
+            if let api::LogStreamEvent::EventLog(log) = event {
+                output::remote_log_event(&log);
+            }
+        }
+    })
+    .await;
+}
+
 async fn do_follow_run(config: Config, run: &Run) -> Result<(), Error> {
-    let enable_ctrl_c = !output::get_output_mode().is_mcp(); // Disable Ctrl+C in MCP mode
+    let enable_ctrl_c = !output::get_output_mode().is_mcp();
     let mut spinner = output::spinner("Waiting for run to start...");
     match wait_for_run_start(&config, &run).await {
         Err(err) => {
@@ -257,35 +303,24 @@ async fn do_follow_run(config: Config, run: &Run) -> Result<(), Error> {
 
             // We do this here, explicitly, to not double-monitor our API via the
             // `wait_for_run_start` function above.
-            let mut run_complete = monitor_run_completion(&config, run);
+            let run_complete = monitor_run_completion(&config, run);
 
             // Now we follow the logs from the run. We can stream them from the cloud to here using
             // the stream_logs API endpoint.
             match api::stream_run_logs(&config, &run.app_name, run.number).await {
-                Ok(mut output) => loop {
-                    let should_exit = tokio::select! {
-                        Some(event) = output.recv() => {
-                            if let api::LogStreamEvent::EventLog(log) = &event {
-                                let ts = dates::format_str(&log.reported_at);
-                                output::log_line(&ts, &log.content, output::LogLineType::Remote);
-                            }
-                            false
-                        },
-                        res = &mut run_complete => {
-                            handle_run_completion(res)?;
-                            true
-                        },
-                        _ = tokio::signal::ctrl_c(), if enable_ctrl_c => {
-                            output::write("Received Ctrl+C, stopping log streaming...\n");
-                            output::write("Note: The run will continue in Tower cloud\n");
-                            output::write(&format!("  See more: {}\n", run.dollar_link));
-                            true
-                        },
-                    };
-                    if should_exit {
-                        break;
+                Ok(log_stream) => {
+                    let completed_run = stream_logs_until_complete(
+                        log_stream,
+                        run_complete,
+                        enable_ctrl_c,
+                        &run.dollar_link,
+                    )
+                    .await?;
+
+                    if let Some(run) = completed_run {
+                        handle_run_completion(Ok(run))?;
                     }
-                },
+                }
                 Err(err) => {
                     output::error(&format!("Failed to stream run logs: {:?}", err));
                     return Err(Error::LogStreamFailed);
@@ -297,9 +332,7 @@ async fn do_follow_run(config: Config, run: &Run) -> Result<(), Error> {
     Ok(())
 }
 
-fn handle_run_completion(
-    res: Result<Run, tokio::sync::oneshot::error::RecvError>,
-) -> Result<(), Error> {
+fn handle_run_completion(res: Result<Run, oneshot::error::RecvError>) -> Result<(), Error> {
     match res {
         Ok(completed_run) => match completed_run.status {
             tower_api::models::run::Status::Errored => {
@@ -540,7 +573,7 @@ async fn monitor_status(app: LocalApp) -> Status {
                             );
                             return tower_runtime::Status::Running; // Return a default status for timeout
                         }
-                        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                        sleep(Duration::from_millis(100)).await;
                         continue;
                     }
                 }
@@ -556,7 +589,7 @@ async fn monitor_status(app: LocalApp) -> Status {
                     output::error("Failed to get app status after timeout");
                     return tower_runtime::Status::Running; // Return a default status for timeout
                 }
-                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                sleep(Duration::from_millis(100)).await;
             }
         }
     }
@@ -585,7 +618,7 @@ async fn wait_for_run_start(config: &Config, run: &Run) -> Result<(), Error> {
             break;
         } else {
             // Wait half a second to to try again.
-            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+            sleep(Duration::from_millis(500)).await;
         }
     }
 
@@ -602,7 +635,7 @@ async fn wait_for_run_completion(config: &Config, run: &Run) -> Result<Run, Erro
             return Ok(res.run);
         } else {
             // Wait half a second to to try again.
-            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+            sleep(Duration::from_millis(500)).await;
         }
     }
 }

tests/integration/features/cli_runs.feature

@@ -31,3 +31,10 @@ Feature: CLI Run Commands
     Then the output should show "Scheduling run..." spinner
     And the output should show "Waiting for run to start..." spinner
     And both spinners should complete successfully
+
+  Scenario: CLI run should show logs that arrive after run completes
+    Given I have a simple hello world application named "app-logs-after-completion"
+    When I run "tower deploy --create" via CLI
+    And I run "tower run" via CLI
+    Then the output should show "First log before run completes"
+    And the output should show "Second log after run completes"

tests/integration/features/steps/mcp_steps.py

@@ -898,3 +898,8 @@ def step_then_receive_workflow_help_stdio(context):
     assert "content" in result
     content = result["content"][0]["text"]
     assert "Tower Application Development Workflow" in content
+
+
+@given('I have a simple hello world application named "{app_name}"')
+def step_create_hello_world_app_named(context, app_name):
+    create_towerfile(context, app_name=app_name)

tests/mock-api-server/main.py

@@ -260,7 +260,11 @@ async def describe_run(name: str, seq: int):
             now_time = datetime.datetime.now()
             elapsed = (now_time - created_time).total_seconds()
 
-            if elapsed > 5:  # After 5 seconds, mark as completed
+            # For logs-after-completion test apps, complete quickly to test log draining
+            # Use 1 second so CLI has time to start streaming before completion
+            completion_threshold = 1.0 if "logs-after-completion" in name else 5.0
+
+            if elapsed > completion_threshold:
                 run_data["status"] = "exited"
                 run_data["status_group"] = "successful"
                 run_data["exit_code"] = 0
@@ -449,58 +453,76 @@ async def refresh_session(refresh_params: Dict[str, Any] = None):
     }
 
 
+NORMAL_LOG_ENTRIES = [
+    (1, "Starting application...", "2025-08-22T12:00:00Z"),
+    (2, "Hello, World!", "2025-08-22T12:00:01Z"),
+    (3, "Application completed successfully", "2025-08-22T12:00:02Z"),
+]
+
+
+def make_log_data(seq: int, line_num: int, content: str, timestamp: str):
+    return {
+        "channel": "program",
+        "content": content,
+        "line_num": line_num,
+        "reported_at": timestamp,
+        "run_id": f"mock-run-{seq}",
+    }
+
+
+def make_log_event(seq: int, line_num: int, content: str, timestamp: str):
+    return f"event: log\ndata: {json.dumps(make_log_data(seq, line_num, content, timestamp))}\n\n"
+
+
 @app.get("/v1/apps/{name}/runs/{seq}/logs")
 async def describe_run_logs(name: str, seq: int):
     """Mock endpoint for getting run logs."""
     if name not in mock_apps_db:
         raise HTTPException(status_code=404, detail=f"App '{name}' not found")
 
-    # Return mock log entries
     return {
         "log_lines": [
-            {"timestamp": "2025-08-22T12:00:00Z", "message": "Starting application..."},
-            {"timestamp": "2025-08-22T12:00:01Z", "message": "Hello, World!"},
-            {
-                "timestamp": "2025-08-22T12:00:02Z",
-                "message": "Application completed successfully",
-            },
+            make_log_data(seq, line_num, content, timestamp)
+            for line_num, content, timestamp in NORMAL_LOG_ENTRIES
         ]
     }
 
 
+async def generate_logs_after_completion_test_stream(seq: int):
+    """Log before run completion, then log after.
+
+    Timeline: Run completes at 1 second, second log sent at 1.5 seconds.
+    """
+    yield make_log_event(
+        seq, 1, "First log before run completes", "2025-08-22T12:00:00Z"
+    )
+    await asyncio.sleep(1.5)
+    yield make_log_event(
+        seq, 2, "Second log after run completes", "2025-08-22T12:00:01Z"
+    )
+
+
+async def generate_normal_log_stream(seq: int):
+    """Normal log stream for regular tests."""
+    for line_num, content, timestamp in NORMAL_LOG_ENTRIES:
+        yield make_log_event(seq, line_num, content, timestamp)
+        await asyncio.sleep(0.1)
+
+
 @app.get("/v1/apps/{name}/runs/{seq}/logs/stream")
 async def stream_run_logs(name: str, seq: int):
     """Mock endpoint for streaming run logs."""
 
     if name not in mock_apps_db:
         raise HTTPException(status_code=404, detail=f"App '{name}' not found")
 
-    async def generate_log_stream():
-        # Simulate streaming logs with proper SSE format for Tower CLI
-        mock_logs = [
-            {"timestamp": "2025-08-22T12:00:00Z", "content": "Starting application..."},
-            {"timestamp": "2025-08-22T12:00:01Z", "content": "Hello, World!"},
-            {
-                "timestamp": "2025-08-22T12:00:02Z",
-                "content": "Application completed successfully",
-            },
-        ]
-
-        for i, log_entry in enumerate(mock_logs):
-            # Format as RunLogLine structure expected by CLI
-            log_data = {
-                "channel": "program",
-                "content": log_entry["content"],
-                "line_num": i + 1,
-                "reported_at": log_entry["timestamp"],
-                "run_id": f"mock-run-{seq}",
-            }
-            # Proper SSE format with event type and data
-            yield f"event: log\ndata: {json.dumps(log_data)}\n\n"
-            await asyncio.sleep(0.1)  # Small delay between logs
+    if "logs-after-completion" in name:
+        stream = generate_logs_after_completion_test_stream(seq)
+    else:
+        stream = generate_normal_log_stream(seq)
 
     return StreamingResponse(
-        generate_log_stream(),
+        stream,
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
     )