use std::path::Path;
use std::sync::OnceLock;
use aion_awl::{CompiledWorkflow, TypeBody};
use aion_package::{
ContentHash, ExtractionLimits, Package, PackageContract, PackageError, SignalContract,
};
use serde_json::Value;
pub const EMBEDDED_ASSISTANT_DOCUMENT: &str = include_str!("../../assistant-embed/assistant.awl");
pub const EMBEDDED_ASSISTANT_FILENAME: &str = "assistant.awl";
pub const OBJECTIVE_INPUT: &str = "objective";
pub const REPO_PATH_INPUT: &str = "repo_path";
pub const CONTINUE_SIGNAL: &str = "assistant_continue";
pub const CONTINUE_MESSAGE_FIELD: &str = "message";
pub const CONTINUE_END_FIELD: &str = "end";
pub const STATUS_QUERY: &str = "assistant_status";
#[must_use]
pub const fn private_task_queue(workflow_type: &str) -> &str {
workflow_type
}
const EMBEDDED_SCHEMA_ROOT: &str = "<embedded-assistant-has-no-schema-directory>";
#[derive(Debug, thiserror::Error)]
pub enum EmbeddedAssistantError {
#[error("the embedded assistant document does not parse: {message}")]
Parse {
message: String,
},
#[error(
"the embedded assistant document imports schema `{path}`, but the binary embeds the \
document alone and has no directory to resolve imports against; declare the type \
inline in the document"
)]
SchemaImport {
path: String,
},
#[error("the embedded assistant document does not compile: {message}")]
Compile {
message: String,
},
#[error("the embedded assistant document could not be packaged: {message}")]
Assemble {
message: String,
},
#[error("the embedded assistant package did not validate: {source}")]
Package {
#[from]
source: PackageError,
},
#[error(
"the embedded assistant document declares no `{name}` input; the session contract in \
crates/aion-server/src/assistant/document.rs names it, so document and contract have \
diverged"
)]
MissingInput {
name: &'static str,
},
#[error(
"the embedded assistant document declares no `{name}` signal; the session contract in \
crates/aion-server/src/assistant/document.rs names it, so document and contract have \
diverged"
)]
MissingSignal {
name: &'static str,
},
#[error(
"the `{signal}` signal payload declares no `{field}` field; the session contract in \
crates/aion-server/src/assistant/document.rs sends it, so document and contract have \
diverged"
)]
MissingSignalField {
signal: &'static str,
field: &'static str,
},
#[error(
"the embedded assistant document declares no `{name}` query; the session contract in \
crates/aion-server/src/assistant/document.rs names it, so document and contract have \
diverged"
)]
MissingQuery {
name: &'static str,
},
#[error(
"the embedded assistant package carries no contract, so its signal payloads cannot be \
read: {message}"
)]
MissingContract {
message: String,
},
#[error(
"the embedded assistant exports workflow type `{workflow_type}`, so its derived private \
queue would be `{default_queue}` — the queue every out-of-box worker comes up on. The \
built-in assistant never claims it (#200); the workflow must be named something else"
)]
QueueWouldBeDefault {
workflow_type: String,
default_queue: &'static str,
},
#[error(
"the embedded assistant document declares no `worker` block, so it claims no task queue; \
the assistant's queue is derived from its own workflow type and the document must \
declare `worker {expected}`"
)]
MissingWorkerBlock {
expected: String,
},
#[error(
"the embedded assistant document declares more than one `worker` block (`{declared}`); \
the built-in assistant serves exactly one derived private queue, `{expected}`"
)]
AmbiguousQueue {
declared: String,
expected: String,
},
#[error(
"the embedded assistant document declares task queue `{declared}`, but the assistant's \
queue is derived from its own workflow type and must be `{expected}`. `default` is the \
out-of-box workers' queue and the built-in assistant never claims it (#200)"
)]
QueueNotDerived {
declared: String,
expected: String,
},
}
#[derive(Debug, Clone)]
pub struct EmbeddedAssistant {
source: &'static str,
package: Package,
workflow_type: String,
task_queue: String,
input_schema: Value,
signals: Vec<SignalContract>,
queries: Vec<String>,
}
impl EmbeddedAssistant {
pub fn load() -> Result<Self, EmbeddedAssistantError> {
Self::from_source(EMBEDDED_ASSISTANT_DOCUMENT)
}
pub fn from_source(source: &'static str) -> Result<Self, EmbeddedAssistantError> {
let document = aion_awl::parse(source).map_err(|error| EmbeddedAssistantError::Parse {
message: error.message,
})?;
for declaration in &document.types {
if let TypeBody::SchemaImport { path, .. } = &declaration.body {
return Err(EmbeddedAssistantError::SchemaImport { path: path.clone() });
}
}
let root = Path::new(EMBEDDED_SCHEMA_ROOT);
let prepared =
aion_awl_package::compile_and_assemble_awl(source, root, EMBEDDED_ASSISTANT_FILENAME)
.map_err(|error| match error {
aion_awl_package::PrepareAwlError::Compile(compile) => {
EmbeddedAssistantError::Compile {
message: compile.to_string(),
}
}
other => EmbeddedAssistantError::Assemble {
message: other.to_string(),
},
})?;
let CompiledWorkflow { input_schema, .. } = prepared.compiled;
let package = Package::load_from_bytes(&prepared.archive, ExtractionLimits::unbounded())?;
let workflow_type = package.manifest().entry_module.clone();
let contract =
package
.contract()
.map_err(|error| EmbeddedAssistantError::MissingContract {
message: error.to_string(),
})?;
let signals = contract.signals.clone();
let task_queue = verified_task_queue(&workflow_type, contract)?;
for name in [OBJECTIVE_INPUT, REPO_PATH_INPUT] {
if !document.inputs.iter().any(|input| input.name == name) {
return Err(EmbeddedAssistantError::MissingInput { name });
}
}
let continuation = signals
.iter()
.find(|signal| signal.name == CONTINUE_SIGNAL)
.ok_or(EmbeddedAssistantError::MissingSignal {
name: CONTINUE_SIGNAL,
})?;
for field in [CONTINUE_MESSAGE_FIELD, CONTINUE_END_FIELD] {
if !schema_declares_property(&continuation.input_schema, field) {
return Err(EmbeddedAssistantError::MissingSignalField {
signal: CONTINUE_SIGNAL,
field,
});
}
}
let queries: Vec<String> = document
.queries
.iter()
.map(|query| query.name.clone())
.collect();
if !queries.iter().any(|name| name == STATUS_QUERY) {
return Err(EmbeddedAssistantError::MissingQuery { name: STATUS_QUERY });
}
Ok(Self {
source,
package,
workflow_type,
task_queue,
input_schema,
signals,
queries,
})
}
#[must_use]
pub const fn package(&self) -> &Package {
&self.package
}
#[must_use]
pub fn workflow_type(&self) -> &str {
&self.workflow_type
}
#[must_use]
pub fn task_queue(&self) -> &str {
&self.task_queue
}
#[must_use]
pub const fn content_hash(&self) -> &ContentHash {
self.package.content_hash()
}
#[must_use]
pub const fn source(&self) -> &'static str {
self.source
}
#[must_use]
pub const fn input_schema(&self) -> &Value {
&self.input_schema
}
#[must_use]
pub fn signals(&self) -> &[SignalContract] {
&self.signals
}
#[must_use]
pub fn queries(&self) -> &[String] {
&self.queries
}
pub fn continuation_schema(&self) -> Result<&Value, EmbeddedAssistantError> {
self.signals
.iter()
.find(|signal| signal.name == CONTINUE_SIGNAL)
.map(|signal| &signal.input_schema)
.ok_or(EmbeddedAssistantError::MissingSignal {
name: CONTINUE_SIGNAL,
})
}
}
fn verified_task_queue(
workflow_type: &str,
contract: &PackageContract,
) -> Result<String, EmbeddedAssistantError> {
let expected = private_task_queue(workflow_type);
if expected == aion_core::DEFAULT_TASK_QUEUE {
return Err(EmbeddedAssistantError::QueueWouldBeDefault {
workflow_type: workflow_type.to_owned(),
default_queue: aion_core::DEFAULT_TASK_QUEUE,
});
}
let declared: Vec<&str> = contract
.workers
.iter()
.map(|worker| worker.task_queue.as_str())
.collect();
let [only] = declared.as_slice() else {
return Err(if declared.is_empty() {
EmbeddedAssistantError::MissingWorkerBlock {
expected: expected.to_owned(),
}
} else {
EmbeddedAssistantError::AmbiguousQueue {
declared: declared.join(", "),
expected: expected.to_owned(),
}
});
};
if *only != expected {
return Err(EmbeddedAssistantError::QueueNotDerived {
declared: (*only).to_owned(),
expected: expected.to_owned(),
});
}
Ok(expected.to_owned())
}
fn schema_declares_property(schema: &Value, property: &str) -> bool {
resolve_local_ref(schema)
.and_then(|resolved| resolved.get("properties"))
.and_then(Value::as_object)
.is_some_and(|properties| properties.contains_key(property))
}
fn resolve_local_ref(schema: &Value) -> Option<&Value> {
let Some(reference) = schema.get("$ref").and_then(Value::as_str) else {
return Some(schema);
};
let name = reference.strip_prefix("#/$defs/")?;
schema.get("$defs")?.get(name)
}
pub fn embedded_assistant() -> Result<&'static EmbeddedAssistant, &'static EmbeddedAssistantError> {
static EMBEDDED: OnceLock<Result<EmbeddedAssistant, EmbeddedAssistantError>> = OnceLock::new();
EMBEDDED.get_or_init(EmbeddedAssistant::load).as_ref()
}
#[cfg(test)]
#[path = "document_tests.rs"]
mod document_tests;