coyote-ai 0.8.3

An all-in-one, batteries included LLM CLI Tool
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

use anyhow::{Result, anyhow};
use chrono::Utc;
use indexmap::IndexMap;
use parking_lot::RwLock;
use std::sync::LazyLock;

type AccessTokenEntry = (String, i64, Option<String>);

static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> =
    LazyLock::new(|| RwLock::new(IndexMap::new()));

pub fn get_access_token(client_name: &str) -> Result<String> {
    ACCESS_TOKENS
        .read()
        .get(client_name)
        .map(|(token, _, _)| token.clone())
        .ok_or_else(|| anyhow!("Invalid access token"))
}

pub fn get_access_token_account_id(client_name: &str) -> Option<String> {
    ACCESS_TOKENS
        .read()
        .get(client_name)
        .and_then(|(_, _, account_id)| account_id.clone())
}

pub fn is_valid_access_token(client_name: &str) -> bool {
    let access_tokens = ACCESS_TOKENS.read();
    let (token, expires_at, _) = match access_tokens.get(client_name) {
        Some(v) => v,
        None => return false,
    };
    !token.is_empty() && Utc::now().timestamp() < *expires_at
}

pub fn set_access_token(
    client_name: &str,
    token: String,
    expires_at: i64,
    account_id: Option<String>,
) {
    let mut access_tokens = ACCESS_TOKENS.write();
    let entry = access_tokens.entry(client_name.to_string()).or_default();
    entry.0 = token;
    entry.1 = expires_at;
    entry.2 = account_id;
}