bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: resolve upstream auth at runtime

Browse Source
This commit is contained in:
a.tolmachev
2026-04-07 01:00:24 +03:00
parent 191e749b14
commit da94d308de
17 changed files with 702 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/mcp-server/src/app.rs
+143 -15
View File
@@ -1,4 +1,5 @@
use std::{
collections::BTreeMap,
convert::Infallible,
sync::Arc,
time::{Duration, Instant},
@@ -19,15 +20,16 @@ use axum::{
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::{
AsyncJobHandle, AsyncJobId, InvocationLevel, InvocationLog, InvocationLogId, InvocationSource,
InvocationStatus, JobStatus, PlatformApiKeyScope, StreamSession, StreamSessionId, StreamStatus,
AsyncJobHandle, AsyncJobId, AuthProfile, InvocationLevel, InvocationLog, InvocationLogId,
InvocationSource, InvocationStatus, JobStatus, PlatformApiKeyScope, SecretId, StreamSession,
StreamSessionId, StreamStatus,
};
use crank_registry::{
CreateAsyncJobRequest, CreateInvocationLogRequest, CreateStreamSessionRequest,
PostgresRegistry, PublishedAgentTool, UpdateAsyncJobStatusRequest,
UpdateStreamSessionStateRequest,
};
use crank_runtime::{RuntimeError, RuntimeExecutor, RuntimeOperation};
use crank_runtime::{ResolvedAuth, RuntimeError, RuntimeExecutor, RuntimeOperation, SecretCrypto};
use futures_util::stream;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
@@ -58,6 +60,7 @@ pub struct AppState {
registry: PostgresRegistry,
catalog: PublishedToolCatalog,
runtime: RuntimeExecutor,
secret_crypto: SecretCrypto,
sessions: SessionStore,
allowed_origins: AllowedOrigins,
}
@@ -118,11 +121,13 @@ pub fn build_app(
registry: PostgresRegistry,
refresh_interval: Duration,
public_base_url: Option<String>,
secret_crypto: SecretCrypto,
) -> Router {
let state = Arc::new(AppState {
registry: registry.clone(),
catalog: PublishedToolCatalog::new(registry, refresh_interval),
runtime: RuntimeExecutor::new(),
secret_crypto,
sessions: SessionStore::new(),
allowed_origins: AllowedOrigins::new(public_base_url),
});
@@ -509,6 +514,87 @@ struct StoredSessionState {
batch_size: usize,
}
async fn resolve_operation_auth(
state: &Arc<AppState>,
workspace_id: &crank_core::WorkspaceId,
execution_config: &crank_core::ExecutionConfig,
) -> Result<Option<ResolvedAuth>, RuntimeError> {
resolve_runtime_auth_for_task(
&state.registry,
&state.secret_crypto,
workspace_id,
execution_config,
)
.await
}
async fn resolve_runtime_auth_for_task(
registry: &PostgresRegistry,
secret_crypto: &SecretCrypto,
workspace_id: &crank_core::WorkspaceId,
execution_config: &crank_core::ExecutionConfig,
) -> Result<Option<ResolvedAuth>, RuntimeError> {
let Some(auth_profile_id) = execution_config.auth_profile_ref.as_ref() else {
return Ok(None);
};
let auth_profile = registry
.get_auth_profile(workspace_id, auth_profile_id)
.await
.map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?
.ok_or_else(|| RuntimeError::MissingAuthProfile {
auth_profile_id: auth_profile_id.as_str().to_owned(),
})?;
resolve_auth_profile(registry, secret_crypto, workspace_id, &auth_profile)
.await
.map(Some)
}
async fn resolve_auth_profile(
registry: &PostgresRegistry,
secret_crypto: &SecretCrypto,
workspace_id: &crank_core::WorkspaceId,
auth_profile: &AuthProfile,
) -> Result<ResolvedAuth, RuntimeError> {
let mut secrets = BTreeMap::new();
let used_at = now_rfc3339();
for secret_id in auth_profile.config.secret_ids() {
let secret = registry
.get_secret(workspace_id, secret_id)
.await
.map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?
.ok_or_else(|| RuntimeError::MissingSecret {
secret_id: secret_id.as_str().to_owned(),
})?;
let version = registry
.get_current_secret_version(workspace_id, secret_id)
.await
.map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?
.ok_or_else(|| RuntimeError::MissingSecretVersion {
secret_id: secret_id.as_str().to_owned(),
version: secret.secret.current_version,
})?;
let plaintext = secret_crypto.decrypt(&version.secret_version.ciphertext)?;
registry
.touch_secret(workspace_id, secret_id, &used_at)
.await
.map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?;
secrets.insert(SecretId::new(secret_id.as_str()), plaintext);
}
ResolvedAuth::from_profile(auth_profile, &secrets)
}
async fn handle_base_tool_call(
state: Arc<AppState>,
session: &SessionState,
@@ -528,11 +614,13 @@ async fn handle_base_tool_call(
.map(|streaming| streaming.mode),
Some(crank_core::ExecutionMode::Window)
);
let resolved_auth =
resolve_operation_auth(&state, &tool.workspace_id, &operation.execution_config).await;
let result = if is_window_mode {
state
let result = match resolved_auth {
Ok(resolved_auth) if is_window_mode => state
.runtime
.execute_window(&operation, &arguments)
.execute_window_with_auth(&operation, &arguments, resolved_auth.as_ref())
.await
.map(|output| {
json!({
@@ -543,9 +631,14 @@ async fn handle_base_tool_call(
"truncated": output.truncated,
"has_more": output.has_more,
})
})
} else {
state.runtime.execute(&operation, &arguments).await
}),
Ok(resolved_auth) => {
state
.runtime
.execute_with_auth(&operation, &arguments, resolved_auth.as_ref())
.await
}
Err(error) => Err(error),
};
match result {
@@ -619,11 +712,25 @@ async fn handle_session_start_call(
);
};
match state
.runtime
.execute_session_seed(&runtime_operation, &arguments)
.await
{
let resolved_auth = resolve_operation_auth(
&state,
&tool.workspace_id,
&runtime_operation.execution_config,
)
.await;
match match resolved_auth {
Ok(resolved_auth) => {
state
.runtime
.execute_session_seed_with_auth(
&runtime_operation,
&arguments,
resolved_auth.as_ref(),
)
.await
}
Err(error) => Err(error),
} {
Ok(seed) => {
let batch_size = streaming.max_items.unwrap_or(10).max(1) as usize;
let preview_count = seed.items.len().min(batch_size);
@@ -983,13 +1090,28 @@ async fn handle_async_job_start_call(
}
let registry = state.registry.clone();
let secret_crypto = state.secret_crypto.clone();
let tool_for_task = tool.clone();
let arguments_for_task = arguments.clone();
let job_id = job.id.clone();
tokio::spawn(async move {
let runtime = RuntimeExecutor::new();
let task_operation = runtime_operation(&tool_for_task);
let result = runtime.execute(&task_operation, &arguments_for_task).await;
let resolved_auth = resolve_runtime_auth_for_task(
&registry,
&secret_crypto,
&tool_for_task.workspace_id,
&task_operation.execution_config,
)
.await;
let result = match resolved_auth {
Ok(resolved_auth) => {
runtime
.execute_with_auth(&task_operation, &arguments_for_task, resolved_auth.as_ref())
.await
}
Err(error) => Err(error),
};
let finished_at = now_rfc3339();
let update_result = match result {
@@ -1732,6 +1854,12 @@ fn runtime_error_code(error: &RuntimeError) -> &'static str {
RuntimeError::UnsupportedExecutionMode { .. } => "streaming_mode_error",
RuntimeError::InvalidPreparedRequest { .. } => "runtime_error",
RuntimeError::InvalidStreamingPayload { .. } => "streaming_payload_error",
RuntimeError::MissingAuthProfile { .. } => "auth_profile_not_found",
RuntimeError::MissingSecret { .. } | RuntimeError::MissingSecretVersion { .. } => {
"secret_not_found"
}
RuntimeError::InvalidAuthSecretValue { .. } => "secret_value_error",
RuntimeError::SecretCrypto { .. } => "secret_crypto_error",
}
}
apps/mcp-server/src/main.rs
+33 -17
View File
@@ -6,6 +6,7 @@ mod session;
use std::{env, net::SocketAddr, time::Duration};
use crank_registry::PostgresRegistry;
use crank_runtime::SecretCrypto;
use tokio::net::TcpListener;
use tracing::info;
@@ -30,7 +31,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.unwrap_or_else(|| Duration::from_secs(5));
let socket_addr: SocketAddr = bind_addr.parse()?;
let registry = PostgresRegistry::connect(&database_url).await?;
let app = build_app(registry, refresh_interval, public_base_url);
let secret_crypto = SecretCrypto::new(&env::var("CRANK_MASTER_KEY")?)?;
let app = build_app(registry, refresh_interval, public_base_url, secret_crypto);
let listener = TcpListener::bind(socket_addr).await?;
info!("mcp-server listening on {}", socket_addr);
@@ -68,6 +70,7 @@ mod tests {
CreateAgentRequest, CreatePlatformApiKeyRequest, ListInvocationLogsQuery, PostgresRegistry,
PublishAgentRequest, PublishRequest,
};
use crank_runtime::SecretCrypto;
use crank_schema::{Schema, SchemaKind};
use futures_util::stream;
use serde_json::{Value, json};
@@ -86,6 +89,19 @@ mod tests {
"default"
}
fn build_test_app(
registry: PostgresRegistry,
refresh_interval: Duration,
public_base_url: Option<String>,
) -> axum::Router {
build_app(
registry,
refresh_interval,
public_base_url,
SecretCrypto::new("test-master-key").unwrap(),
)
}
#[tokio::test]
async fn initializes_lists_and_calls_published_tool_via_mcp() {
let registry = test_registry().await;
@@ -114,7 +130,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry.clone(),
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -220,7 +236,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -284,7 +300,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -326,7 +342,7 @@ mod tests {
publish_agent_with_bindings(&registry, "sales-init", vec![]).await;
let api_key =
create_platform_api_key(&registry, "mcp-init", &[PlatformApiKeyScope::Read]).await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -388,7 +404,7 @@ mod tests {
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -438,7 +454,7 @@ mod tests {
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -489,7 +505,7 @@ mod tests {
&[PlatformApiKeyScope::Read],
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -518,7 +534,7 @@ mod tests {
&[PlatformApiKeyScope::Read],
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -562,7 +578,7 @@ mod tests {
&[PlatformApiKeyScope::Read],
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -589,7 +605,7 @@ mod tests {
async fn refreshes_published_tools_without_restart() {
let registry = test_registry().await;
let upstream_base_url = spawn_upstream_server().await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry.clone(),
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -700,7 +716,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -747,7 +763,7 @@ mod tests {
async fn rejects_initialize_without_platform_api_key() {
let registry = test_registry().await;
publish_agent_with_bindings(&registry, "sales-auth", vec![]).await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -796,7 +812,7 @@ mod tests {
let api_key =
create_platform_api_key(&registry, "mcp-read", &[PlatformApiKeyScope::Read]).await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -857,7 +873,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry.clone(),
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -1001,7 +1017,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry.clone(),
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
@@ -1148,7 +1164,7 @@ mod tests {
)
.await;
let base_url = spawn_mcp_server(build_app(
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),