Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 6 additions & 57 deletions aw-client-rust/src/classes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,27 +2,29 @@
//!
//! Taken from default classes in aw-webui

use log::warn;
use rand::Rng;
use serde::{Deserialize, Serialize};

use super::blocking::AwClient as ActivityWatchClient;
use serde_json;

pub type CategoryId = Vec<String>;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CategorySpec {
#[serde(rename = "type")]
pub spec_type: String,
#[serde(default)]
pub regex: String,
#[serde(default)]
pub ignore_case: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassSetting {
#[serde(default)]
pub id: Option<i32>,
pub name: Vec<String>,
pub rule: CategorySpec,
#[serde(default)]
pub data: Option<serde_json::Value>,
}

/// Returns the default categorization classes
Expand Down Expand Up @@ -140,56 +142,3 @@ pub fn default_classes() -> Vec<(CategoryId, CategorySpec)> {
),
]
}

/// Get classes from server-side settings using default localhost:5600.
/// Might throw an error if not set yet, in which case we use the default classes as a fallback.
pub fn get_classes() -> Vec<(CategoryId, CategorySpec)> {
get_classes_from_server("localhost", 5600)
}

/// Get classes from server-side settings with custom host and port.
/// Might throw an error if not set yet, in which case we use the default classes as a fallback.
pub fn get_classes_from_server(host: &str, port: u16) -> Vec<(CategoryId, CategorySpec)> {
let mut rng = rand::rng();
let random_int = rng.random_range(0..10001);
let client_id = format!("get-setting-{}", random_int);

// Create a client with a random ID, similar to the Python implementation
let awc = match ActivityWatchClient::new(host, port, &client_id) {
Ok(client) => client,
Err(_) => {
warn!(
"Failed to create ActivityWatch client for {}:{}, using default classes",
host, port
);
return default_classes();
}
};

awc.get_setting("classes")
.map(|setting_value| {
// Try to deserialize the setting into Vec<ClassSetting>
if setting_value.is_null() {
return default_classes();
}

let class_settings: Vec<ClassSetting> = serde_json::from_value(setting_value)
.unwrap_or_else(|_| {
warn!("Failed to deserialize classes setting, using default classes");
return vec![];
});

// Convert ClassSetting to (CategoryId, CategorySpec) format
class_settings
.into_iter()
.map(|class| (class.name, class.rule))
.collect()
})
.unwrap_or_else(|_| {
warn!(
"Failed to get classes from server {}:{}, using default classes as fallback",
host, port
);
default_classes()
})
}
2 changes: 1 addition & 1 deletion aw-client-rust/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ impl std::fmt::Debug for AwClient {
}

fn get_hostname() -> String {
return gethostname::gethostname().to_string_lossy().to_string();
gethostname::gethostname().to_string_lossy().to_string()
}

impl AwClient {
Expand Down
151 changes: 67 additions & 84 deletions aw-client-rust/src/queries.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,8 @@
//! };
//!
//! // Automatically fetches classes from localhost:5600
//! let query = QueryParams::Desktop(params.clone()).canonical_events_with_classes();
//! let query = QueryParams::Desktop(params.clone()).canonical_events();
//!
//! // Or from a custom server
//! let query = QueryParams::Desktop(params)
//! .canonical_events_with_classes_from_server("localhost", 2345);
//! ```

use crate::classes::{CategoryId, CategorySpec};
Expand Down Expand Up @@ -78,8 +75,6 @@ pub static BROWSER_APPNAMES: phf::Map<&'static str, &'static [&'static str]> = p
"vivaldi" => &["Vivaldi-stable", "Vivaldi-snapshot", "vivaldi.exe"],
};

pub const DEFAULT_LIMIT: u32 = 100;

/// Type alias for categorization classes
pub type ClassRule = (CategoryId, CategorySpec);

Expand Down Expand Up @@ -135,54 +130,46 @@ impl QueryParams {
QueryParams::Android(params) => build_android_canonical_events(params),
}
}
}

/// Build canonical events query string with automatic class fetching if not provided
pub fn canonical_events_with_classes(&self) -> String {
self.canonical_events_with_classes_from_server("localhost", 5600)
}
/// Helper function to serialize classes in the format expected by the categorize function
/// This version builds the query string directly without JSON serialization to avoid double-escaping
fn serialize_classes(classes: &[ClassRule]) -> String {
let mut parts = Vec::new();

for (category_id, category_spec) in classes {
// Build category array string manually: ["Work", "Programming"]
let category_str = format!(
"[{}]",
category_id
.iter()
.map(|s| format!("\"{}\"", s))
.collect::<Vec<_>>()
.join(", ")
);

/// Build canonical events query string with automatic class fetching from custom server
pub fn canonical_events_with_classes_from_server(&self, host: &str, port: u16) -> String {
match self {
QueryParams::Desktop(params) => {
let mut params_with_classes = params.clone();
if params_with_classes.base.classes.is_empty() {
params_with_classes.base.classes =
crate::classes::get_classes_from_server(host, port);
}
build_desktop_canonical_events(&params_with_classes)
}
QueryParams::Android(params) => {
let mut params_with_classes = params.clone();
if params_with_classes.base.classes.is_empty() {
params_with_classes.base.classes =
crate::classes::get_classes_from_server(host, port);
}
build_android_canonical_events(&params_with_classes)
}
// Build spec object manually to avoid JSON escaping regex patterns
let mut spec_parts = Vec::new();
spec_parts.push(format!("\"type\": \"{}\"", category_spec.spec_type));

// Only include regex for non-"none" types, and use raw pattern without escaping
if category_spec.spec_type != "none" {
spec_parts.push(format!("\"regex\": \"{}\"", category_spec.regex));
}

// Always include ignore_case field
spec_parts.push(format!("\"ignore_case\": {}", category_spec.ignore_case));

let spec_str = format!("{{{}}}", spec_parts.join(", "));

// Build the tuple [category, spec]
parts.push(format!("[{}, {}]", category_str, spec_str));
}
}

/// Helper function to serialize classes in the format expected by the categorize function
fn serialize_classes(classes: &[ClassRule]) -> String {
// Convert Vec<(CategoryId, CategorySpec)> to the JSON format expected by categorize
let serialized_classes: Vec<(Vec<String>, serde_json::Value)> = classes
.iter()
.map(|(category_id, category_spec)| {
let spec_json = serde_json::json!({
"type": category_spec.spec_type,
"regex": category_spec.regex,
"ignore_case": category_spec.ignore_case
});
(category_id.clone(), spec_json)
})
.collect();

serde_json::to_string(&serialized_classes).unwrap_or_else(|_| "[]".to_string())
format!("[{}]", parts.join(", "))
}

fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
pub fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
let mut query = Vec::new();

// Fetch window events
Expand All @@ -195,7 +182,7 @@ fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
if params.base.filter_afk {
query.push(format!(
"not_afk = flood(query_bucket(find_bucket(\"{}\")));
not_afk = filter_keyvals(not_afk, \"status\", [\"not-afk\"])",
not_afk = filter_keyvals(not_afk, \"status\", [\"not-afk\"])",
escape_doublequote(&params.bid_afk)
));
}
Expand All @@ -207,7 +194,7 @@ fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
if params.base.include_audible {
query.push(
"audible_events = filter_keyvals(browser_events, \"audible\", [true]);
not_afk = period_union(not_afk, audible_events)"
not_afk = period_union(not_afk, audible_events)"
.to_string(),
);
}
Expand All @@ -221,7 +208,7 @@ fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
// Add categorization if classes specified
if !params.base.classes.is_empty() {
query.push(format!(
"events = categorize(events, {})",
"events = categorize(events, {});",
serialize_classes(&params.base.classes)
));
}
Expand All @@ -237,7 +224,7 @@ fn build_desktop_canonical_events(params: &DesktopQueryParams) -> String {
query.join(";\n")
}

fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
pub fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
let mut query = Vec::new();

// Fetch app events
Expand All @@ -252,7 +239,7 @@ fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
// Add categorization if classes specified
if !params.base.classes.is_empty() {
query.push(format!(
"events = categorize(events, {})",
"events = categorize(events, {});",
serialize_classes(&params.base.classes)
));
}
Expand All @@ -268,62 +255,58 @@ fn build_android_canonical_events(params: &AndroidQueryParams) -> String {
query.join(";\n")
}

fn build_browser_events(params: &DesktopQueryParams) -> String {
let mut query = String::from("browser_events = [];\n");
pub fn build_browser_events(params: &DesktopQueryParams) -> String {
let mut query = String::from("browser_events = [];");

for browser_bucket in &params.base.bid_browsers {
for (browser_name, app_names) in BROWSER_APPNAMES.entries() {
if browser_bucket.contains(browser_name) {
query.push_str(&format!(
"events_{0} = flood(query_bucket(\"{1}\"));
window_{0} = filter_keyvals(events, \"app\", {2});
events_{0} = filter_period_intersect(events_{0}, window_{0});
events_{0} = split_url_events(events_{0});
browser_events = concat(browser_events, events_{0});
browser_events = sort_by_timestamp(browser_events);\n",
"
events_{0} = flood(query_bucket(\"{1}\"));
window_{0} = filter_keyvals(events, \"app\", {2});
events_{0} = filter_period_intersect(events_{0}, window_{0});
events_{0} = split_url_events(events_{0});
browser_events = concat(browser_events, events_{0});
browser_events = sort_by_timestamp(browser_events)",
browser_name,
escape_doublequote(browser_bucket),
serde_json::to_string(app_names).unwrap()
));
}
}
}

query
}

/// Build a full desktop query
/// Build a full desktop query using default localhost:5600 configuration
pub fn full_desktop_query(params: &DesktopQueryParams) -> String {
let mut query = QueryParams::Desktop(params.clone()).canonical_events_with_classes();
let mut query = QueryParams::Desktop(params.clone()).canonical_events();

// Add basic event aggregations
query.push_str(&format!(
"
query.push_str(
&"
title_events = sort_by_duration(merge_events_by_keys(events, [\"app\", \"title\"]));
app_events = sort_by_duration(merge_events_by_keys(title_events, [\"app\"]));
cat_events = sort_by_duration(merge_events_by_keys(events, [\"$category\"]));
app_events = limit_events(app_events, {});
title_events = limit_events(title_events, {});
duration = sum_durations(events);
",
DEFAULT_LIMIT, DEFAULT_LIMIT
));
"
.to_string(),
);

// Add browser-specific query parts if browser buckets exist
if !params.base.bid_browsers.is_empty() {
query.push_str(&format!(
"
query.push_str(
&"
browser_events = split_url_events(browser_events);
browser_urls = merge_events_by_keys(browser_events, [\"url\"]);
browser_urls = sort_by_duration(browser_urls);
browser_urls = limit_events(browser_urls, {});
browser_domains = merge_events_by_keys(browser_events, [\"$domain\"]);
browser_domains = sort_by_duration(browser_domains);
browser_domains = limit_events(browser_domains, {});
browser_duration = sum_durations(browser_events);
",
DEFAULT_LIMIT, DEFAULT_LIMIT
));
"
.to_string(),
);
} else {
query.push_str(
"
Expand Down Expand Up @@ -414,17 +397,17 @@ mod tests {
assert!(serialized.contains("Programming"));
assert!(serialized.contains("Google Docs"));
assert!(serialized.contains("GitHub|vim"));
assert!(serialized.contains("\"type\":\"regex\""));
assert!(serialized.contains("\"ignore_case\":false"));
assert!(serialized.contains("\"ignore_case\":true"));
assert!(serialized.contains("\"type\": \"regex\""));
assert!(serialized.contains("\"ignore_case\": false"));
assert!(serialized.contains("\"ignore_case\": true"));
}

#[test]
fn test_canonical_events_with_empty_classes() {
let params = DesktopQueryParams {
base: QueryParamsBase {
bid_browsers: vec![],
classes: vec![], // Empty classes - should trigger server fetch
classes: vec![],
filter_classes: vec![],
filter_afk: true,
include_audible: true,
Expand All @@ -434,9 +417,9 @@ mod tests {
};

let query_params = QueryParams::Desktop(params);
let query = query_params.canonical_events_with_classes();
let query = query_params.canonical_events();

// Should contain basic query structure even if server fetch fails
// Should contain basic query structure
assert!(query.contains("events = flood"));
assert!(query.contains("test-window"));
}
Expand Down Expand Up @@ -465,7 +448,7 @@ mod tests {
};

let query_params = QueryParams::Desktop(params);
let query = query_params.canonical_events_with_classes();
let query = query_params.canonical_events();

// Should contain categorization
assert!(query.contains("events = categorize"));
Expand Down
Loading