[codex] Support API key spend sorting

crates/api/src/routes/workspaces.rs

@@ -154,6 +154,48 @@ impl From<WorkspaceOrderDirection> for services::workspace::WorkspaceOrderDirect
     }
 }
 
+#[derive(Debug, Deserialize)]
+pub struct ListApiKeysParams {
+    #[serde(default = "crate::routes::common::default_limit")]
+    pub limit: i64,
+    #[serde(default)]
+    pub offset: i64,
+    pub order_by: Option<ApiKeyOrderBy>,
+    pub order_direction: Option<ApiKeyOrderDirection>,
+}
+
+#[derive(Debug, Deserialize, ToSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum ApiKeyOrderBy {
+    CreatedAt,
+    Usage,
+}
+
+impl From<ApiKeyOrderBy> for services::workspace::ApiKeyOrderBy {
+    fn from(value: ApiKeyOrderBy) -> Self {
+        match value {
+            ApiKeyOrderBy::CreatedAt => Self::CreatedAt,
+            ApiKeyOrderBy::Usage => Self::Usage,
+        }
+    }
+}
+
+#[derive(Debug, Deserialize, ToSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum ApiKeyOrderDirection {
+    Asc,
+    Desc,
+}
+
+impl From<ApiKeyOrderDirection> for services::workspace::ApiKeyOrderDirection {
+    fn from(value: ApiKeyOrderDirection) -> Self {
+        match value {
+            ApiKeyOrderDirection::Asc => Self::Asc,
+            ApiKeyOrderDirection::Desc => Self::Desc,
+        }
+    }
+}
+
 // ============================================
 // Workspace Management Routes
 // ============================================
@@ -772,7 +814,9 @@ pub async fn create_workspace_api_key(
     params(
         ("workspace_id" = Uuid, Path, description = "Workspace ID"),
         ("limit" = Option<i64>, Query, description = "Maximum number of results (default: 20)"),
-        ("offset" = Option<i64>, Query, description = "Number of results to skip (default: 0)")
+        ("offset" = Option<i64>, Query, description = "Number of results to skip (default: 0)"),
+        ("order_by" = Option<ApiKeyOrderBy>, Query, description = "Field to order by: created_at or usage"),
+        ("order_direction" = Option<ApiKeyOrderDirection>, Query, description = "Sort direction: asc or desc")
     ),
     responses(
         (status = 200, description = "Paginated list of workspace API keys", body = ListApiKeysResponse),
@@ -789,7 +833,7 @@ pub async fn list_workspace_api_keys(
     State(app_state): State<AppState>,
     Extension(user): Extension<AuthenticatedUser>,
     Path(workspace_id): Path<Uuid>,
-    Query(params): Query<ListParams>,
+    Query(params): Query<ListApiKeysParams>,
 ) -> Result<Json<ListApiKeysResponse>, (StatusCode, Json<ErrorResponse>)> {
     debug!(
         "Listing API keys for workspace: {} by user: {} (limit: {}, offset: {})",
@@ -837,9 +881,18 @@ pub async fn list_workspace_api_keys(
     };
 
     // Use workspace service to list workspace API keys with pagination and usage data
+    let order_by = params.order_by.map(Into::into);
+    let order_direction = params.order_direction.map(Into::into);
     match app_state
         .workspace_service
-        .list_api_keys_paginated(workspace_id_typed, user_id, params.limit, params.offset)
+        .list_api_keys_paginated(
+            workspace_id_typed,
+            user_id,
+            params.limit,
+            params.offset,
+            order_by,
+            order_direction,
+        )
         .await
     {
         Ok(api_keys) => {

crates/api/tests/e2e_all/api_keys.rs

@@ -88,6 +88,157 @@ async fn test_list_workspace_api_keys() {
     println!("Listed {} API keys successfully", api_keys.len());
 }
 
+#[tokio::test]
+async fn test_list_workspace_api_keys_orders_by_usage() {
+    let (server, database) = setup_test_server_with_database().await;
+    let org = create_org(&server).await;
+    let workspaces = list_workspaces(&server, org.id.clone()).await;
+    let workspace = workspaces.first().unwrap();
+    let model_name = setup_qwen_model(&server).await;
+    let service = get_or_create_web_search_service(&server).await;
+
+    let unused_key =
+        create_api_key_in_workspace(&server, workspace.id.clone(), "Unused Key".to_string()).await;
+    let inference_spend_key = create_api_key_in_workspace(
+        &server,
+        workspace.id.clone(),
+        "Inference Spend Key".to_string(),
+    )
+    .await;
+    let service_spend_key = create_api_key_in_workspace(
+        &server,
+        workspace.id.clone(),
+        "Service Spend Key".to_string(),
+    )
+    .await;
+
+    let organization_id = uuid::Uuid::parse_str(&org.id).unwrap();
+    let workspace_id = uuid::Uuid::parse_str(&workspace.id).unwrap();
+    let unused_key_id = uuid::Uuid::parse_str(&unused_key.id).unwrap();
+    let inference_spend_key_id = uuid::Uuid::parse_str(&inference_spend_key.id).unwrap();
+    let service_spend_key_id = uuid::Uuid::parse_str(&service_spend_key.id).unwrap();
+    let service_id = service.id;
+    let client = database.pool().get().await.unwrap();
+    let model_id: uuid::Uuid = client
+        .query_one(
+            "SELECT id FROM models WHERE model_name = $1",
+            &[&model_name],
+        )
+        .await
+        .unwrap()
+        .get(0);
+
+    for (api_key_id, total_cost) in [
+        (service_spend_key_id, 100_000_000_i64),
+        (inference_spend_key_id, 300_000_000_i64),
+    ] {
+        client
+            .execute(
+                r#"
+                INSERT INTO organization_usage_log (
+                    id, organization_id, workspace_id, api_key_id,
+                    model_id, model_name, input_tokens, output_tokens,
+                    total_tokens, input_cost, output_cost, total_cost,
+                    inference_type, created_at
+                ) VALUES ($1, $2, $3, $4, $5, $6, 10, 10, 20, 1, 1, $7,
+                          'chat_completion', NOW())
+                "#,
+                &[
+                    &uuid::Uuid::new_v4(),
+                    &organization_id,
+                    &workspace_id,
+                    &api_key_id,
+                    &model_id,
+                    &model_name,
+                    &total_cost,
+                ],
+            )
+            .await
+            .unwrap();
+    }
+
+    client
+        .execute(
+            r#"
+            INSERT INTO organization_service_usage_log (
+                id, organization_id, workspace_id, api_key_id,
+                service_id, quantity, total_cost, inference_id, created_at
+            ) VALUES ($1, $2, $3, $4, $5, 1, 400000000, NULL, NOW())
+            "#,
+            &[
+                &uuid::Uuid::new_v4(),
+                &organization_id,
+                &workspace_id,
+                &service_spend_key_id,
+                &service_id,
+            ],
+        )
+        .await
+        .unwrap();
+
+    client
+        .execute(
+            "UPDATE api_keys SET created_at = NOW() + INTERVAL '3 hours' WHERE id = $1",
+            &[&service_spend_key_id],
+        )
+        .await
+        .unwrap();
+    client
+        .execute(
+            "UPDATE api_keys SET created_at = NOW() + INTERVAL '2 hours' WHERE id = $1",
+            &[&inference_spend_key_id],
+        )
+        .await
+        .unwrap();
+    client
+        .execute(
+            "UPDATE api_keys SET created_at = NOW() + INTERVAL '1 hour' WHERE id = $1",
+            &[&unused_key_id],
+        )
+        .await
+        .unwrap();
+
+    let default_response = server
+        .get(format!("/v1/workspaces/{}/api-keys?limit=3", workspace.id).as_str())
+        .add_header("Authorization", format!("Bearer {}", get_session_id()))
+        .await;
+    assert_eq!(default_response.status_code(), 200);
+    let default_list = default_response.json::<api::models::ListApiKeysResponse>();
+    assert_eq!(default_list.api_keys[0].id, service_spend_key.id);
+    assert_eq!(default_list.api_keys[1].id, inference_spend_key.id);
+    assert_eq!(default_list.api_keys[2].id, unused_key.id);
+
+    let desc_response = server
+        .get(
+            format!(
+                "/v1/workspaces/{}/api-keys?limit=2&order_by=usage&order_direction=desc",
+                workspace.id
+            )
+            .as_str(),
+        )
+        .add_header("Authorization", format!("Bearer {}", get_session_id()))
+        .await;
+    assert_eq!(desc_response.status_code(), 200);
+    let desc_list = desc_response.json::<api::models::ListApiKeysResponse>();
+    assert_eq!(desc_list.api_keys[0].id, service_spend_key.id);
+    assert_eq!(desc_list.api_keys[1].id, inference_spend_key.id);
+
+    let asc_response = server
+        .get(
+            format!(
+                "/v1/workspaces/{}/api-keys?limit=2&order_by=usage&order_direction=asc",
+                workspace.id
+            )
+            .as_str(),
+        )
+        .add_header("Authorization", format!("Bearer {}", get_session_id()))
+        .await;
+    assert_eq!(asc_response.status_code(), 200);
+    let asc_list = asc_response.json::<api::models::ListApiKeysResponse>();
+    assert_eq!(asc_list.api_keys[0].id, unused_key.id);
+    assert_eq!(asc_list.api_keys[1].id, inference_spend_key.id);
+}
+
 #[tokio::test]
 async fn test_api_key_prevents_duplicate_names_in_workspace() {
     let server = setup_test_server().await;

crates/database/src/migrations/sql/V0067__add_service_usage_api_key_index.sql

@@ -0,0 +1,2 @@
+CREATE INDEX IF NOT EXISTS idx_org_service_usage_workspace_api_key
+    ON organization_service_usage_log(workspace_id, api_key_id);

crates/database/src/repositories/api_key.rs

@@ -6,7 +6,7 @@ use anyhow::{Context, Result};
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use services::common::{extract_api_key_prefix, generate_api_key, hash_api_key, RepositoryError};
-use services::workspace::ports::CreateApiKeyRequest;
+use services::workspace::ports::{ApiKeyOrderBy, ApiKeyOrderDirection, CreateApiKeyRequest};
 use tracing::debug;
 use uuid::Uuid;
 
@@ -237,7 +237,25 @@ impl ApiKeyRepository {
         workspace_id: Uuid,
         limit: i64,
         offset: i64,
+        order_by: Option<ApiKeyOrderBy>,
+        order_direction: Option<ApiKeyOrderDirection>,
     ) -> Result<Vec<ApiKey>, RepositoryError> {
+        let order_by = order_by.unwrap_or(ApiKeyOrderBy::CreatedAt);
+        let order_direction = order_direction.unwrap_or(ApiKeyOrderDirection::Desc);
+
+        let order_by_column = match order_by {
+            ApiKeyOrderBy::CreatedAt => "ak.created_at",
+            ApiKeyOrderBy::Usage => "usage",
+        };
+        let order_dir = match order_direction {
+            ApiKeyOrderDirection::Asc => "ASC",
+            ApiKeyOrderDirection::Desc => "DESC",
+        };
+        let tie_breaker = match order_by {
+            ApiKeyOrderBy::CreatedAt => ", ak.id ASC",
+            ApiKeyOrderBy::Usage => ", ak.created_at DESC, ak.id ASC",
+        };
+
         let rows = retry_db!("list_api_keys_by_workspace_paginated", {
             let client = self
                 .pool
@@ -248,7 +266,8 @@ impl ApiKeyRepository {
 
             client
                 .query(
-                    r#"
+                    &format!(
+                        r#"
                 SELECT 
                     ak.id,
                     ak.key_hash,
@@ -262,14 +281,28 @@ impl ApiKeyRepository {
                     ak.is_active,
                     ak.deleted_at,
                     ak.spend_limit,
-                    COALESCE(SUM(usg.total_cost), 0)::BIGINT as usage
+                    (
+                        COALESCE(inference_usage.total_cost, 0)
+                        + COALESCE(service_usage.total_cost, 0)
+                    )::BIGINT as usage
                 FROM api_keys ak
-                LEFT JOIN organization_usage_log usg ON ak.id = usg.api_key_id
+                LEFT JOIN (
+                    SELECT api_key_id, COALESCE(SUM(total_cost), 0)::BIGINT AS total_cost
+                    FROM organization_usage_log
+                    WHERE workspace_id = $1
+                    GROUP BY api_key_id
+                ) inference_usage ON ak.id = inference_usage.api_key_id
+                LEFT JOIN (
+                    SELECT api_key_id, COALESCE(SUM(total_cost), 0)::BIGINT AS total_cost
+                    FROM organization_service_usage_log
+                    WHERE workspace_id = $1
+                    GROUP BY api_key_id
+                ) service_usage ON ak.id = service_usage.api_key_id
                 WHERE ak.workspace_id = $1 AND ak.deleted_at IS NULL
-                GROUP BY ak.id
-                ORDER BY ak.created_at DESC
+                ORDER BY {order_by_column} {order_dir}{tie_breaker}
                 LIMIT $2 OFFSET $3
-                "#,
+                "#
+                    ),
                     &[&workspace_id, &limit, &offset],
                 )
                 .await
@@ -620,9 +653,11 @@ impl services::workspace::ports::ApiKeyRepository for ApiKeyRepository {
         workspace_id: services::workspace::WorkspaceId,
         limit: i64,
         offset: i64,
+        order_by: Option<services::workspace::ApiKeyOrderBy>,
+        order_direction: Option<services::workspace::ApiKeyOrderDirection>,
     ) -> Result<Vec<services::workspace::ApiKey>, RepositoryError> {
         let api_keys = self
-            .list_by_workspace_paginated(workspace_id.0, limit, offset)
+            .list_by_workspace_paginated(workspace_id.0, limit, offset, order_by, order_direction)
             .await?;
         Ok(api_keys
             .into_iter()

crates/services/src/auth/mod.rs

@@ -567,8 +567,9 @@ mod tests {
         UpdateOrganizationMemberRequest, UpdateOrganizationRequest,
     };
     use crate::workspace::{
-        ApiKey, ApiKeyId, ApiKeyRepository, CreateApiKeyRequest, Workspace, WorkspaceId,
-        WorkspaceOrderBy, WorkspaceOrderDirection, WorkspaceRepository,
+        ApiKey, ApiKeyId, ApiKeyOrderBy, ApiKeyOrderDirection, ApiKeyRepository,
+        CreateApiKeyRequest, Workspace, WorkspaceId, WorkspaceOrderBy, WorkspaceOrderDirection,
+        WorkspaceRepository,
     };
     use bloomfilter::Bloom;
     use chrono::Utc;
@@ -735,6 +736,8 @@ mod tests {
             _: WorkspaceId,
             _: i64,
             _: i64,
+            _: Option<ApiKeyOrderBy>,
+            _: Option<ApiKeyOrderDirection>,
         ) -> Result<Vec<ApiKey>, RepositoryError> {
             unimplemented!()
         }

crates/services/src/workspace/mod.rs

@@ -253,14 +253,16 @@ impl WorkspaceServiceTrait for WorkspaceServiceImpl {
         requester_id: UserId,
         limit: i64,
         offset: i64,
+        order_by: Option<ApiKeyOrderBy>,
+        order_direction: Option<ApiKeyOrderDirection>,
     ) -> Result<Vec<ApiKey>, WorkspaceError> {
         // Check permissions
         self.check_workspace_permission(workspace_id.clone(), requester_id)
             .await?;
 
         // List API keys with pagination (repository now includes usage data via JOIN)
         self.api_key_repository
-            .list_by_workspace_paginated(workspace_id, limit, offset)
+            .list_by_workspace_paginated(workspace_id, limit, offset, order_by, order_direction)
             .await
             .map_err(|e| {
                 WorkspaceError::InternalError(format!(