Skip to content
Open
Show file tree
Hide file tree
Changes from 3 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
9 changes: 9 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -493,12 +493,21 @@ Environment variables override config file values. For CI/CD or one-off use:
| Variable | Default | Description |
|----------|---------|-------------|
| `TOKSCALE_NATIVE_TIMEOUT_MS` | `300000` (5 min) | Overrides `nativeTimeoutMs` config |
| `TOKSCALE_API_URL` | `https://tokscale.ai` | Overrides the API base URL used by `login` and `submit` |
| `TOKSCALE_SOURCE_ID` | auto-generated per machine | Overrides the stable source ID attached to `submit` payloads |
| `TOKSCALE_SOURCE_NAME` | `CLI on <hostname>` | Overrides the human-readable source name attached to `submit` payloads |

```bash
# Example: Increase timeout for very large datasets
TOKSCALE_NATIVE_TIMEOUT_MS=600000 tokscale graph --output data.json

# Example: submit two local test sources without using two machines
TOKSCALE_SOURCE_ID=machine-a tokscale submit
TOKSCALE_SOURCE_ID=machine-b TOKSCALE_SOURCE_NAME="Work Laptop" tokscale submit
```

> `TOKSCALE_SOURCE_NAME` is stored per `sourceId`. A later submit with the same `sourceId` and a new non-empty `TOKSCALE_SOURCE_NAME` updates the stored display name for that source.

> **Note**: For persistent changes, prefer setting `nativeTimeoutMs` in `~/.config/tokscale/settings.json`. Environment variables are best for one-off overrides or CI/CD.

### Headless Mode
Expand Down
122 changes: 121 additions & 1 deletion crates/tokscale-cli/src/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,9 @@ use serde::{Deserialize, Serialize};
use std::fs;
use std::io::IsTerminal;
use std::io::Write;
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use std::thread;
use std::time::Duration;

fn home_dir() -> Result<PathBuf> {
dirs::home_dir().context("Could not determine home directory")
Expand Down Expand Up @@ -53,6 +55,14 @@ fn get_credentials_path() -> Result<PathBuf> {
Ok(home_dir()?.join(".config/tokscale/credentials.json"))
}

fn get_source_id_path() -> Result<PathBuf> {
Ok(home_dir()?.join(".config/tokscale/source-id"))
}

fn get_source_id_lock_path() -> Result<PathBuf> {
Ok(home_dir()?.join(".config/tokscale/source-id.lock"))
}

fn ensure_config_dir() -> Result<()> {
let config_dir = home_dir()?.join(".config/tokscale");

Expand Down Expand Up @@ -125,6 +135,116 @@ fn get_device_name() -> String {
format!("CLI on {}", hostname)
}

fn read_source_id(path: &Path) -> Option<String> {
let content = fs::read_to_string(path).ok()?;
let trimmed = content.trim();
if trimmed.is_empty() {
return None;
}
Some(trimmed.to_string())
}

struct SourceIdLock {
path: PathBuf,
}

impl Drop for SourceIdLock {
fn drop(&mut self) {
let _ = fs::remove_file(&self.path);
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Outdated
}
}

fn acquire_source_id_lock() -> Result<SourceIdLock> {
ensure_config_dir()?;
let lock_path = get_source_id_lock_path()?;

for _ in 0..100 {
match fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&lock_path)
{
Ok(_) => return Ok(SourceIdLock { path: lock_path }),
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
if let Ok(metadata) = fs::metadata(&lock_path) {
if let Ok(modified_at) = metadata.modified() {
if modified_at.elapsed().unwrap_or_default() > Duration::from_secs(10) {
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Outdated
let _ = fs::remove_file(&lock_path);
continue;
}
}
}
thread::sleep(Duration::from_millis(25));
}
Err(err) => return Err(err.into()),
}
}

anyhow::bail!(
"Could not acquire source ID lock after 100 retries (~2500ms)"
);
}

fn write_source_id(path: &Path, source_id: &str) -> Result<()> {
let temp_path = path.with_extension(format!("tmp-{}", std::process::id()));

#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;

let mut file = fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.mode(0o600)
.open(&temp_path)?;
file.write_all(source_id.as_bytes())?;
file.write_all(b"\n")?;
}

#[cfg(not(unix))]
{
fs::write(&temp_path, format!("{source_id}\n"))?;
}

fs::rename(&temp_path, path)?;
Ok(())
}

pub fn get_submit_source_id() -> Result<String> {
if let Some(source_id) = std::env::var_os("TOKSCALE_SOURCE_ID") {
let trimmed = source_id.to_string_lossy().trim().to_string();
if !trimmed.is_empty() {
return Ok(trimmed);
}
}

ensure_config_dir()?;
let path = get_source_id_path()?;

if let Some(existing) = read_source_id(&path) {
return Ok(existing);
}

let _lock = acquire_source_id_lock()?;
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.

if let Some(existing) = read_source_id(&path) {
return Ok(existing);
}

let source_id = uuid::Uuid::new_v4().to_string();
write_source_id(&path, &source_id)?;
Ok(source_id)
}

pub fn get_submit_source_name() -> Option<String> {
std::env::var("TOKSCALE_SOURCE_NAME")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.or_else(|| Some(get_device_name()))
}

#[cfg(target_os = "linux")]
fn has_non_empty_env_var(name: &str) -> bool {
std::env::var_os(name).is_some_and(|value| !value.is_empty())
Expand Down
18 changes: 15 additions & 3 deletions crates/tokscale-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2580,6 +2580,10 @@ struct TsDataSummary {
struct TsExportMeta {
generated_at: String,
version: String,
#[serde(skip_serializing_if = "Option::is_none")]
source_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
source_name: Option<String>,
date_range: DateRange,
}

Expand All @@ -2592,11 +2596,17 @@ struct TsTokenContributionData {
contributions: Vec<TsDailyContribution>,
}

fn to_ts_token_contribution_data(graph: &tokscale_core::GraphResult) -> TsTokenContributionData {
fn to_ts_token_contribution_data(
graph: &tokscale_core::GraphResult,
source_id: Option<String>,
source_name: Option<String>,
) -> TsTokenContributionData {
TsTokenContributionData {
meta: TsExportMeta {
generated_at: graph.meta.generated_at.clone(),
version: graph.meta.version.clone(),
source_id,
source_name,
date_range: DateRange {
start: graph.meta.date_range_start.clone(),
end: graph.meta.date_range_end.clone(),
Expand Down Expand Up @@ -2874,7 +2884,7 @@ fn run_graph_command(
.map_err(|e| anyhow::anyhow!(e))?;

let processing_time_ms = start.elapsed().as_millis() as u32;
let output_data = to_ts_token_contribution_data(&graph_result);
let output_data = to_ts_token_contribution_data(&graph_result, None, None);
let json_output = serde_json::to_string_pretty(&output_data)?;

if let Some(output_path) = output {
Expand Down Expand Up @@ -3118,7 +3128,9 @@ fn run_submit_command(

let api_url = auth::get_api_base_url();

let submit_payload = to_ts_token_contribution_data(&graph_result);
let source_id = auth::get_submit_source_id()?;
let source_name = auth::get_submit_source_name();
let submit_payload = to_ts_token_contribution_data(&graph_result, Some(source_id), source_name);

let response = rt.block_on(async {
reqwest::Client::new()
Expand Down
Loading