pub mod tools;
use crate::error::AppError;
use crate::registrations::AppState;
use axum::http::request::Parts;
use forklaunch_axum::agent::{jwt_policy_from_env, mount_agent};
use forklaunch_axum::native::{
ContractRouter, EndpointError, EndpointPolicy, PolicyCapabilities, RouteContract,
SessionPolicy, Shutdown, StreamOptions,
};
use forklaunch_axum::session::AuthenticatedSession;
use forklaunch_core::agent::{Agent, AgentBuilder{{#has_servers}}, McpServer, env, env_bearer{{/has_servers}}{{#has_services}}, service_tools{{/has_services}}};
use std::error::Error;
use std::time::Duration;
pub const INSTRUCTIONS: &str = {{{instructions}}};
#[must_use]
pub fn builder() -> AgentBuilder {
tools::all()
.into_iter()
.fold(
Agent::builder("{{name}}").instructions(INSTRUCTIONS),
AgentBuilder::tool,
)
{{#services}}
.service(service_tools("{{service}}"))
{{/services}}
{{#servers}}
.mcp_server(
McpServer::new("{{server}}", env("{{url_setting}}"))
.header("authorization", env_bearer("{{token_setting}}")),
)
{{/servers}}
}
#[derive(Clone)]
pub struct AgentRuntime {
pub agent: Agent,
pub auth: SessionPolicy,
}
impl AgentRuntime {
#[must_use]
pub fn new(agent: Agent, auth: SessionPolicy) -> Self {
Self { agent, auth }
}
pub fn from_env() -> Result<Self, AppError> {
let agent = builder()
.build()
.map_err(|error| AppError::InvalidConfig(error.to_string()))?;
let auth = jwt_policy_from_env().map_err(|error| AppError::InvalidConfig(error.to_string()))?;
Ok(Self::new(agent, auth))
}
pub async fn check(&self) -> Result<(), Box<dyn Error>> {
self.agent.check().await?;
Ok(())
}
}
#[derive(Clone)]
pub struct AgentAuth;
impl EndpointPolicy<AppState> for AgentAuth {
const CAPABILITIES: PolicyCapabilities = PolicyCapabilities::Authorized;
type Context = AuthenticatedSession;
async fn authorize(
&self,
parts: &mut Parts,
state: &AppState,
route: &RouteContract,
) -> Result<AuthenticatedSession, EndpointError> {
state.agent.auth.authorize(parts, state, route).await
}
}
pub fn mount(router: ContractRouter<AppState>) -> Result<ContractRouter<AppState>, String> {
let options = StreamOptions::new(
Shutdown::new(),
512 * 1024,
Duration::from_secs(30),
Duration::from_secs(10),
)?;
mount_agent(router, |state: &AppState| state.agent.agent.clone(), AgentAuth, options)
}