attune-system/attune
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

http_request action working nicely

Browse Source
This commit is contained in:
David Culbreth
2026-02-09 23:21:23 -06:00
parent e31ecb781b
commit 966a5af188
18 changed files with 720 additions and 395 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/worker/src/runtime/local.rs
View File

@@ -123,7 +123,7 @@ impl Runtime for LocalRuntime {
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::{ParameterDelivery, ParameterFormat};
use crate::runtime::{OutputFormat, ParameterDelivery, ParameterFormat};
use std::collections::HashMap;
#[tokio::test]

16
crates/worker/src/runtime/python.rs
View File

@@ -339,13 +339,15 @@ if __name__ == '__main__':
None
}
OutputFormat::Json => {
// Try to parse last line of stdout as JSON
stdout_result
.content
.trim()
.lines()
.last()
.and_then(|line| serde_json::from_str(line).ok())
// Try to parse full stdout as JSON first (handles multi-line JSON),
// then fall back to last line only (for scripts that log before output)
let trimmed = stdout_result.content.trim();
serde_json::from_str(trimmed).ok().or_else(|| {
trimmed
.lines()
.last()
.and_then(|line| serde_json::from_str(line).ok())
})
}
OutputFormat::Yaml => {
// Try to parse stdout as YAML

109
crates/worker/src/runtime/shell.rs
View File

@@ -208,13 +208,15 @@ impl ShellRuntime {
None
}
OutputFormat::Json => {
// Try to parse last line of stdout as JSON
stdout_result
.content
.trim()
.lines()
.last()
.and_then(|line| serde_json::from_str(line).ok())
// Try to parse full stdout as JSON first (handles multi-line JSON),
// then fall back to last line only (for scripts that log before output)
let trimmed = stdout_result.content.trim();
serde_json::from_str(trimmed).ok().or_else(|| {
trimmed
.lines()
.last()
.and_then(|line| serde_json::from_str(line).ok())
})
}
OutputFormat::Yaml => {
// Try to parse stdout as YAML
@@ -823,4 +825,97 @@ echo '{"id": 3, "name": "Charlie"}'
assert_eq!(items[2]["id"], 3);
assert_eq!(items[2]["name"], "Charlie");
}
#[tokio::test]
async fn test_shell_runtime_multiline_json_output() {
// Regression test: scripts that embed pretty-printed JSON (e.g., http_request.sh
// embedding a multi-line response body in its "json" field) produce multi-line
// stdout. The parser must handle this by trying to parse the full stdout as JSON
// before falling back to last-line parsing.
let runtime = ShellRuntime::new();
let context = ExecutionContext {
execution_id: 7,
action_ref: "test.multiline_json".to_string(),
parameters: HashMap::new(),
env: HashMap::new(),
secrets: HashMap::new(),
timeout: Some(10),
working_dir: None,
entry_point: "shell".to_string(),
code: Some(
r#"
# Simulate http_request.sh output with embedded pretty-printed JSON
printf '{"status_code":200,"body":"hello","json":{\n "args": {\n "hello": "world"\n },\n "url": "https://example.com"\n},"success":true}\n'
"#
.to_string(),
),
code_path: None,
runtime_name: Some("shell".to_string()),
max_stdout_bytes: 10 * 1024 * 1024,
max_stderr_bytes: 10 * 1024 * 1024,
parameter_delivery: attune_common::models::ParameterDelivery::default(),
parameter_format: attune_common::models::ParameterFormat::default(),
output_format: attune_common::models::OutputFormat::Json,
};
let result = runtime.execute(context).await.unwrap();
assert!(result.is_success());
assert_eq!(result.exit_code, 0);
// Verify result was parsed (not stored as raw stdout)
let parsed = result
.result
.expect("Multi-line JSON should be parsed successfully");
assert_eq!(parsed["status_code"], 200);
assert_eq!(parsed["success"], true);
assert_eq!(parsed["json"]["args"]["hello"], "world");
// stdout should be empty when result is successfully parsed
assert!(
result.stdout.is_empty(),
"stdout should be empty when result is parsed, got: {}",
result.stdout
);
}
#[tokio::test]
async fn test_shell_runtime_json_with_log_prefix() {
// Verify last-line fallback still works: scripts that log to stdout
// before the final JSON line should still parse correctly.
let runtime = ShellRuntime::new();
let context = ExecutionContext {
execution_id: 8,
action_ref: "test.json_with_logs".to_string(),
parameters: HashMap::new(),
env: HashMap::new(),
secrets: HashMap::new(),
timeout: Some(10),
working_dir: None,
entry_point: "shell".to_string(),
code: Some(
r#"
echo "Starting action..."
echo "Processing data..."
echo '{"result": "success", "count": 42}'
"#
.to_string(),
),
code_path: None,
runtime_name: Some("shell".to_string()),
max_stdout_bytes: 10 * 1024 * 1024,
max_stderr_bytes: 10 * 1024 * 1024,
parameter_delivery: attune_common::models::ParameterDelivery::default(),
parameter_format: attune_common::models::ParameterFormat::default(),
output_format: attune_common::models::OutputFormat::Json,
};
let result = runtime.execute(context).await.unwrap();
assert!(result.is_success());
let parsed = result.result.expect("Last-line JSON should be parsed");
assert_eq!(parsed["result"], "success");
assert_eq!(parsed["count"], 42);
}
}

136
crates/worker/src/service.rs
View File

@@ -19,6 +19,7 @@ use sqlx::PgPool;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tracing::{error, info, warn};
use crate::artifacts::ArtifactManager;
@@ -51,6 +52,7 @@ pub struct WorkerService {
mq_connection: Arc<Connection>,
publisher: Arc<Publisher>,
consumer: Option<Arc<Consumer>>,
consumer_handle: Option<JoinHandle<()>>,
worker_id: Option<i64>,
}
@@ -266,6 +268,7 @@ impl WorkerService {
mq_connection: Arc::new(mq_connection),
publisher: Arc::new(publisher),
consumer: None,
consumer_handle: None,
worker_id: None,
})
}
@@ -305,25 +308,35 @@ impl WorkerService {
Ok(())
}
/// Stop the worker service
/// Stop the worker service gracefully
///
/// Shutdown order (mirrors sensor service pattern):
/// 1. Deregister worker (mark inactive to stop receiving new work)
/// 2. Stop heartbeat
/// 3. Wait for in-flight tasks with timeout
/// 4. Close MQ connection
/// 5. Close DB connection
pub async fn stop(&mut self) -> Result<()> {
info!("Stopping Worker Service - initiating graceful shutdown");
// Mark worker as inactive first to stop receiving new tasks
// 1. Mark worker as inactive first to stop receiving new tasks
// Use if-let instead of ? so shutdown continues even if DB call fails
{
let reg = self.registration.read().await;
info!("Marking worker as inactive to stop receiving new tasks");
reg.deregister().await?;
if let Err(e) = reg.deregister().await {
error!("Failed to deregister worker: {}", e);
}
}
// Stop heartbeat
// 2. Stop heartbeat
info!("Stopping heartbeat updates");
self.heartbeat.stop().await;
// Wait a bit for heartbeat to stop
// Wait a bit for heartbeat loop to notice the flag
tokio::time::sleep(Duration::from_millis(100)).await;
// Wait for in-flight tasks to complete (with timeout)
// 3. Wait for in-flight tasks to complete (with timeout)
let shutdown_timeout = self
.config
.worker
@@ -342,6 +355,23 @@ impl WorkerService {
Err(_) => warn!("Shutdown timeout reached - some tasks may have been interrupted"),
}
// 4. Abort consumer task and close message queue connection
if let Some(handle) = self.consumer_handle.take() {
info!("Stopping consumer task...");
handle.abort();
// Wait briefly for the task to finish
let _ = handle.await;
}
info!("Closing message queue connection...");
if let Err(e) = self.mq_connection.close().await {
warn!("Error closing message queue: {}", e);
}
// 5. Close database connection
info!("Closing database connection...");
self.db_pool.close().await;
info!("Worker Service stopped");
Ok(())
@@ -364,6 +394,9 @@ impl WorkerService {
}
/// Start consuming execution.scheduled messages
///
/// Spawns the consumer loop as a background task so that `start()` returns
/// immediately, allowing the caller to set up signal handlers.
async fn start_execution_consumer(&mut self) -> Result<()> {
let worker_id = self
.worker_id
@@ -375,48 +408,63 @@ impl WorkerService {
info!("Starting consumer for worker queue: {}", queue_name);
// Create consumer
let consumer = Consumer::new(
&self.mq_connection,
ConsumerConfig {
queue: queue_name.clone(),
tag: format!("worker-{}", worker_id),
prefetch_count: 10,
auto_ack: false,
exclusive: false,
},
)
.await
.map_err(|e| Error::Internal(format!("Failed to create consumer: {}", e)))?;
info!("Consumer started for queue: {}", queue_name);
info!("Message queue consumer initialized");
// Clone Arc references for the handler
let executor = self.executor.clone();
let publisher = self.publisher.clone();
let db_pool = self.db_pool.clone();
// Consume messages with handler
consumer
.consume_with_handler(
move |envelope: MessageEnvelope<ExecutionScheduledPayload>| {
let executor = executor.clone();
let publisher = publisher.clone();
let db_pool = db_pool.clone();
async move {
Self::handle_execution_scheduled(executor, publisher, db_pool, envelope)
.await
.map_err(|e| format!("Execution handler error: {}", e).into())
}
let consumer = Arc::new(
Consumer::new(
&self.mq_connection,
ConsumerConfig {
queue: queue_name.clone(),
tag: format!("worker-{}", worker_id),
prefetch_count: 10,
auto_ack: false,
exclusive: false,
},
)
.await
.map_err(|e| Error::Internal(format!("Failed to start consumer: {}", e)))?;
.map_err(|e| Error::Internal(format!("Failed to create consumer: {}", e)))?,
);
// Store consumer reference
self.consumer = Some(Arc::new(consumer));
info!("Consumer created for queue: {}", queue_name);
// Clone Arc references for the spawned task
let executor = self.executor.clone();
let publisher = self.publisher.clone();
let db_pool = self.db_pool.clone();
let consumer_for_task = consumer.clone();
let queue_name_for_log = queue_name.clone();
// Spawn the consumer loop as a background task so start() can return
let handle = tokio::spawn(async move {
info!("Consumer loop started for queue '{}'", queue_name_for_log);
let result = consumer_for_task
.consume_with_handler(
move |envelope: MessageEnvelope<ExecutionScheduledPayload>| {
let executor = executor.clone();
let publisher = publisher.clone();
let db_pool = db_pool.clone();
async move {
Self::handle_execution_scheduled(executor, publisher, db_pool, envelope)
.await
.map_err(|e| format!("Execution handler error: {}", e).into())
}
},
)
.await;
match result {
Ok(()) => info!("Consumer loop for queue '{}' ended", queue_name_for_log),
Err(e) => error!(
"Consumer loop for queue '{}' failed: {}",
queue_name_for_log, e
),
}
});
// Store consumer reference and task handle
self.consumer = Some(consumer);
self.consumer_handle = Some(handle);
info!("Message queue consumer initialized");
Ok(())
}