mirror of
https://github.com/harivansh-afk/sandbox-agent.git
synced 2026-04-15 09:01:17 +00:00
fix(pi-agent): add credential extraction and env injection for multi-provider support
Co-authored-by: Nathan <github@nathanflurry.com>
This commit is contained in:
parent
1c381c552a
commit
27dfa37b5a
4 changed files with 125 additions and 9 deletions
|
|
@ -45,6 +45,7 @@ OAuth tokens (used when OAuth extraction is enabled):
|
||||||
| Amp | Anthropic |
|
| Amp | Anthropic |
|
||||||
| Codex | OpenAI |
|
| Codex | OpenAI |
|
||||||
| OpenCode | Anthropic or OpenAI |
|
| OpenCode | Anthropic or OpenAI |
|
||||||
|
| Pi | Any supported provider API key (Anthropic, OpenAI, Gemini, etc.) |
|
||||||
| Mock | None |
|
| Mock | None |
|
||||||
|
|
||||||
## Error handling behavior
|
## Error handling behavior
|
||||||
|
|
|
||||||
|
|
@ -8,6 +8,9 @@ use acp_http_adapter::process::{AdapterError, AdapterRuntime, PostOutcome};
|
||||||
use acp_http_adapter::registry::LaunchSpec;
|
use acp_http_adapter::registry::LaunchSpec;
|
||||||
use axum::response::sse::Event;
|
use axum::response::sse::Event;
|
||||||
use futures::Stream;
|
use futures::Stream;
|
||||||
|
use sandbox_agent_agent_credentials::{
|
||||||
|
extract_all_credentials, AuthType, CredentialExtractionOptions, ExtractedCredentials,
|
||||||
|
};
|
||||||
use sandbox_agent_agent_management::agents::{AgentId, AgentManager, InstallOptions};
|
use sandbox_agent_agent_management::agents::{AgentId, AgentManager, InstallOptions};
|
||||||
use sandbox_agent_error::SandboxError;
|
use sandbox_agent_error::SandboxError;
|
||||||
use sandbox_agent_opencode_adapter::{AcpDispatch, AcpDispatchResult, AcpPayloadStream};
|
use sandbox_agent_opencode_adapter::{AcpDispatch, AcpDispatchResult, AcpPayloadStream};
|
||||||
|
|
@ -303,6 +306,14 @@ impl AcpProxyRuntime {
|
||||||
message: err.to_string(),
|
message: err.to_string(),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
|
let credentials = tokio::task::spawn_blocking(move || {
|
||||||
|
extract_all_credentials(&CredentialExtractionOptions::new())
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|err| SandboxError::StreamError {
|
||||||
|
message: format!("failed to resolve credentials: {err}"),
|
||||||
|
})?;
|
||||||
|
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
server_id = server_id,
|
server_id = server_id,
|
||||||
agent = agent.as_str(),
|
agent = agent.as_str(),
|
||||||
|
|
@ -312,11 +323,14 @@ impl AcpProxyRuntime {
|
||||||
"create_instance: launch spec resolved, spawning"
|
"create_instance: launch spec resolved, spawning"
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let mut launch_env = launch.env;
|
||||||
|
merge_credentials_env(&mut launch_env, &credentials);
|
||||||
|
|
||||||
let runtime = AdapterRuntime::start(
|
let runtime = AdapterRuntime::start(
|
||||||
LaunchSpec {
|
LaunchSpec {
|
||||||
program: launch.program,
|
program: launch.program,
|
||||||
args: launch.args,
|
args: launch.args,
|
||||||
env: launch.env,
|
env: launch_env,
|
||||||
},
|
},
|
||||||
self.inner.request_timeout,
|
self.inner.request_timeout,
|
||||||
)
|
)
|
||||||
|
|
@ -488,6 +502,41 @@ fn annotate_agent_error(agent: AgentId, mut value: Value) -> Value {
|
||||||
value
|
value
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn merge_credentials_env(env: &mut HashMap<String, String>, credentials: &ExtractedCredentials) {
|
||||||
|
if let Some(cred) = &credentials.anthropic {
|
||||||
|
if cred.auth_type == AuthType::ApiKey {
|
||||||
|
insert_env_if_missing(env, "ANTHROPIC_API_KEY", &cred.api_key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(cred) = &credentials.openai {
|
||||||
|
if cred.auth_type == AuthType::ApiKey {
|
||||||
|
insert_env_if_missing(env, "OPENAI_API_KEY", &cred.api_key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for (provider, cred) in &credentials.other {
|
||||||
|
if cred.auth_type != AuthType::ApiKey {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let key = format!("{}_API_KEY", provider.to_uppercase().replace('-', "_"));
|
||||||
|
insert_env_if_missing(env, &key, &cred.api_key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn insert_env_if_missing(env: &mut HashMap<String, String>, key: &str, value: &str) {
|
||||||
|
if env.contains_key(key) || env_has_value(key) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
env.insert(key.to_string(), value.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
fn env_has_value(key: &str) -> bool {
|
||||||
|
std::env::var(key)
|
||||||
|
.map(|value| !value.trim().is_empty())
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
fn duration_from_env_ms(key: &str, default: Duration) -> Duration {
|
fn duration_from_env_ms(key: &str, default: Duration) -> Duration {
|
||||||
match std::env::var(key) {
|
match std::env::var(key) {
|
||||||
Ok(raw) => raw
|
Ok(raw) => raw
|
||||||
|
|
|
||||||
|
|
@ -17,7 +17,7 @@ use sandbox_agent_agent_management::agents::{
|
||||||
AgentId, AgentManager, InstallOptions, InstallResult, InstallSource, InstalledArtifactKind,
|
AgentId, AgentManager, InstallOptions, InstallResult, InstallSource, InstalledArtifactKind,
|
||||||
};
|
};
|
||||||
use sandbox_agent_agent_management::credentials::{
|
use sandbox_agent_agent_management::credentials::{
|
||||||
extract_all_credentials, CredentialExtractionOptions,
|
extract_all_credentials, AuthType, CredentialExtractionOptions,
|
||||||
};
|
};
|
||||||
use sandbox_agent_error::{ErrorType, ProblemDetails, SandboxError};
|
use sandbox_agent_error::{ErrorType, ProblemDetails, SandboxError};
|
||||||
use sandbox_agent_opencode_adapter::{build_opencode_router, OpenCodeAdapterConfig};
|
use sandbox_agent_opencode_adapter::{build_opencode_router, OpenCodeAdapterConfig};
|
||||||
|
|
@ -427,6 +427,19 @@ async fn get_v1_agents(
|
||||||
|
|
||||||
let has_anthropic = credentials.anthropic.is_some();
|
let has_anthropic = credentials.anthropic.is_some();
|
||||||
let has_openai = credentials.openai.is_some();
|
let has_openai = credentials.openai.is_some();
|
||||||
|
let has_other = !credentials.other.is_empty();
|
||||||
|
let has_any_api_key = credentials
|
||||||
|
.anthropic
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|cred| cred.auth_type == AuthType::ApiKey)
|
||||||
|
|| credentials
|
||||||
|
.openai
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|cred| cred.auth_type == AuthType::ApiKey)
|
||||||
|
|| credentials
|
||||||
|
.other
|
||||||
|
.values()
|
||||||
|
.any(|cred| cred.auth_type == AuthType::ApiKey);
|
||||||
|
|
||||||
let instances = state.acp_proxy().list_instances().await;
|
let instances = state.acp_proxy().list_instances().await;
|
||||||
let mut active_by_agent = HashMap::<AgentId, Vec<i64>>::new();
|
let mut active_by_agent = HashMap::<AgentId, Vec<i64>>::new();
|
||||||
|
|
@ -444,7 +457,13 @@ async fn get_v1_agents(
|
||||||
for agent_id in AgentId::all().iter().copied() {
|
for agent_id in AgentId::all().iter().copied() {
|
||||||
let capabilities = agent_capabilities_for(agent_id);
|
let capabilities = agent_capabilities_for(agent_id);
|
||||||
let installed = state.agent_manager().is_installed(agent_id);
|
let installed = state.agent_manager().is_installed(agent_id);
|
||||||
let credentials_available = credentials_available_for(agent_id, has_anthropic, has_openai);
|
let credentials_available = credentials_available_for(
|
||||||
|
agent_id,
|
||||||
|
has_anthropic,
|
||||||
|
has_openai,
|
||||||
|
has_other,
|
||||||
|
has_any_api_key,
|
||||||
|
);
|
||||||
|
|
||||||
let server_status = active_by_agent.get(&agent_id).map(|created_times| {
|
let server_status = active_by_agent.get(&agent_id).map(|created_times| {
|
||||||
let uptime_ms = created_times
|
let uptime_ms = created_times
|
||||||
|
|
@ -569,6 +588,19 @@ async fn get_v1_agent(
|
||||||
|
|
||||||
let has_anthropic = credentials.anthropic.is_some();
|
let has_anthropic = credentials.anthropic.is_some();
|
||||||
let has_openai = credentials.openai.is_some();
|
let has_openai = credentials.openai.is_some();
|
||||||
|
let has_other = !credentials.other.is_empty();
|
||||||
|
let has_any_api_key = credentials
|
||||||
|
.anthropic
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|cred| cred.auth_type == AuthType::ApiKey)
|
||||||
|
|| credentials
|
||||||
|
.openai
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|cred| cred.auth_type == AuthType::ApiKey)
|
||||||
|
|| credentials
|
||||||
|
.other
|
||||||
|
.values()
|
||||||
|
.any(|cred| cred.auth_type == AuthType::ApiKey);
|
||||||
|
|
||||||
let instances = state.acp_proxy().list_instances().await;
|
let instances = state.acp_proxy().list_instances().await;
|
||||||
let created_times: Vec<i64> = instances
|
let created_times: Vec<i64> = instances
|
||||||
|
|
@ -579,7 +611,13 @@ async fn get_v1_agent(
|
||||||
|
|
||||||
let capabilities = agent_capabilities_for(agent_id);
|
let capabilities = agent_capabilities_for(agent_id);
|
||||||
let installed = state.agent_manager().is_installed(agent_id);
|
let installed = state.agent_manager().is_installed(agent_id);
|
||||||
let credentials_available = credentials_available_for(agent_id, has_anthropic, has_openai);
|
let credentials_available = credentials_available_for(
|
||||||
|
agent_id,
|
||||||
|
has_anthropic,
|
||||||
|
has_openai,
|
||||||
|
has_other,
|
||||||
|
has_any_api_key,
|
||||||
|
);
|
||||||
|
|
||||||
let server_status = if created_times.is_empty() {
|
let server_status = if created_times.is_empty() {
|
||||||
None
|
None
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,7 @@
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use sandbox_agent_agent_management::credentials::{
|
||||||
|
extract_all_credentials, AuthType, CredentialExtractionOptions,
|
||||||
|
};
|
||||||
|
|
||||||
pub(super) async fn not_found() -> Response {
|
pub(super) async fn not_found() -> Response {
|
||||||
let problem = ProblemDetails {
|
let problem = ProblemDetails {
|
||||||
|
|
@ -48,12 +51,15 @@ pub(super) fn credentials_available_for(
|
||||||
agent: AgentId,
|
agent: AgentId,
|
||||||
has_anthropic: bool,
|
has_anthropic: bool,
|
||||||
has_openai: bool,
|
has_openai: bool,
|
||||||
|
has_other: bool,
|
||||||
|
has_any_api_key: bool,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
match agent {
|
match agent {
|
||||||
AgentId::Claude | AgentId::Amp => has_anthropic,
|
AgentId::Claude | AgentId::Amp => has_anthropic,
|
||||||
AgentId::Codex => has_openai,
|
AgentId::Codex => has_openai,
|
||||||
AgentId::Opencode => has_anthropic || has_openai,
|
AgentId::Opencode => has_anthropic || has_openai || has_other,
|
||||||
AgentId::Pi | AgentId::Cursor => true,
|
AgentId::Pi => has_any_api_key,
|
||||||
|
AgentId::Cursor => true,
|
||||||
AgentId::Mock => true,
|
AgentId::Mock => true,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -524,8 +530,22 @@ pub(super) fn build_provider_payload_for_opencode(_state: &Arc<AppState>) -> Val
|
||||||
AgentId::Cursor,
|
AgentId::Cursor,
|
||||||
];
|
];
|
||||||
|
|
||||||
let has_anthropic = std::env::var("ANTHROPIC_API_KEY").is_ok();
|
let credentials = extract_all_credentials(&CredentialExtractionOptions::new());
|
||||||
let has_openai = std::env::var("OPENAI_API_KEY").is_ok();
|
let has_anthropic = credentials.anthropic.is_some();
|
||||||
|
let has_openai = credentials.openai.is_some();
|
||||||
|
let has_other = !credentials.other.is_empty();
|
||||||
|
let has_any_api_key = credentials
|
||||||
|
.anthropic
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|cred| cred.auth_type == AuthType::ApiKey)
|
||||||
|
|| credentials
|
||||||
|
.openai
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|cred| cred.auth_type == AuthType::ApiKey)
|
||||||
|
|| credentials
|
||||||
|
.other
|
||||||
|
.values()
|
||||||
|
.any(|cred| cred.auth_type == AuthType::ApiKey);
|
||||||
|
|
||||||
let mut all_providers = Vec::new();
|
let mut all_providers = Vec::new();
|
||||||
let mut defaults = serde_json::Map::new();
|
let mut defaults = serde_json::Map::new();
|
||||||
|
|
@ -579,7 +599,15 @@ pub(super) fn build_provider_payload_for_opencode(_state: &Arc<AppState>) -> Val
|
||||||
|
|
||||||
defaults.insert(agent_str.to_string(), json!(current_value));
|
defaults.insert(agent_str.to_string(), json!(current_value));
|
||||||
|
|
||||||
if agent == AgentId::Mock || credentials_available_for(agent, has_anthropic, has_openai) {
|
if agent == AgentId::Mock
|
||||||
|
|| credentials_available_for(
|
||||||
|
agent,
|
||||||
|
has_anthropic,
|
||||||
|
has_openai,
|
||||||
|
has_other,
|
||||||
|
has_any_api_key,
|
||||||
|
)
|
||||||
|
{
|
||||||
connected.push(json!(agent_str));
|
connected.push(json!(agent_str));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue