Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ members = [
"crates/aionui-process",
"crates/aionui-session",
"crates/aionui-session-message",
"crates/aionui-skill-runtime",
"crates/aionui-auth",
"crates/aionui-system",
"crates/aionui-file",
Expand Down Expand Up @@ -45,6 +46,7 @@ aionui-runtime = { path = "crates/aionui-runtime" }
aionui-process = { path = "crates/aionui-process" }
aionui-session = { path = "crates/aionui-session" }
aionui-session-message = { path = "crates/aionui-session-message" }
aionui-skill-runtime = { path = "crates/aionui-skill-runtime" }
aionui-auth = { path = "crates/aionui-auth" }
aionui-system = { path = "crates/aionui-system" }
aionui-file = { path = "crates/aionui-file" }
Expand Down
195 changes: 172 additions & 23 deletions crates/aionui-ai-agent/src/capability/first_message_injector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@

use std::sync::Arc;

use aionui_api_types::SkillDeliveryMode;

use crate::capability::skill_manager::{AcpSkillManager, prepare_first_message_with_skills_index};

/// Configuration for the first-message injector.
Expand All @@ -17,25 +19,29 @@ pub struct InjectionConfig<'a> {
pub preset_context: Option<&'a str>,
/// Resolved skill names (snapshot from `conversation.extra.skills`).
pub skills: &'a [String],
/// True iff the agent's native CLI reads skills from the workspace
/// without needing prompt injection. Derived by callers from
/// `AcpBackend::native_skills_dirs().is_some()` for ACP.
pub native_skill_support: bool,
/// How this vendor receives skills.
///
/// This replaced a `native_skill_support: bool` derived from
/// `native_skills_dirs.is_some()`. That signal conflated two different
/// things — "declares a workspace skills directory" and "can discover
/// skills natively" — so a vendor with a declared directory but no working
/// discovery got LIGHT injection and therefore no skills at all.
pub delivery_mode: SkillDeliveryMode,
}

/// Produce the content string to send as the first ACP prompt.
/// Produce the content string to send as the first prompt.
///
/// - If `native_skill_support`: **light mode** — only `preset_context`
/// prepended as an `[Assistant Rules]` block (if present). The native CLI
/// handles skill discovery via workspace links.
/// - Else: **heavy mode** — `preset_context` + resolved skills index
/// injected via `prepare_first_message_with_skills_index`.
/// Two states, from three modes:
/// * `Argv` / `Protocol` — LIGHT: only `preset_context`. The CLI owns skill
/// discovery, and injecting an index would duplicate the name+description
/// that a plugin-registering CLI already adds always-on.
/// * `Injected` — the skills index plus the dual-channel instructions.
pub async fn inject_first_message_prefix(
content: &str,
manager: &Arc<AcpSkillManager>,
config: InjectionConfig<'_>,
) -> String {
if config.native_skill_support {
if is_light_mode(&config.delivery_mode) {
return match config.preset_context {
Some(ctx) if !ctx.is_empty() => {
format!("[Assistant Rules]\n{ctx}\n[/Assistant Rules]\n\n{content}")
Expand All @@ -52,6 +58,31 @@ pub async fn inject_first_message_prefix(
prepare_first_message_with_skills_index(content, &skills, config.preset_context)
}

/// The same block [`inject_first_message_prefix`] would prepend, WITHOUT the
/// user's content appended — `None` when there is nothing to inject.
///
/// For backends that carry the prefix separately from the turn's message rather
/// than concatenating it once: agy re-invokes its CLI per turn, so its rules
/// belong on the first invocation only. Sharing this function with the
/// concatenating path is what keeps the two from drifting into different wording.
pub async fn compose_injected_prefix(manager: &Arc<AcpSkillManager>, config: InjectionConfig<'_>) -> Option<String> {
// A sentinel the caller can split on: `prepare_first_message_with_skills_index`
// owns the block's exact layout, so re-deriving it here would be a second
// copy of that layout.
const SENTINEL: &str = "\u{0}AIONUI_CONTENT\u{0}";
let composed = inject_first_message_prefix(SENTINEL, manager, config).await;
match composed.strip_suffix(SENTINEL) {
// Nothing was prepended: the content came back untouched.
None => None,
Some(prefix) if prefix.trim().is_empty() => None,
Some(prefix) => Some(prefix.trim_end().to_owned()),
}
}

fn is_light_mode(mode: &SkillDeliveryMode) -> bool {
matches!(mode, SkillDeliveryMode::Argv | SkillDeliveryMode::Protocol)
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -82,8 +113,126 @@ mod tests {
}
}

fn skill_corpus(base: &std::path::Path) {
let auto = base.join("auto-inject");
std::fs::create_dir_all(auto.join("cron")).unwrap();
std::fs::write(
auto.join("cron").join("SKILL.md"),
"---\nname: cron\ndescription: Schedule stuff\n---\nBody.",
)
.unwrap();
}

/// BOTH layer-1 modes must be light. `Protocol` used to fall to the heavy
/// branch because the old signal was a single bool: codex would then have
/// received a duplicate index on top of what its CLI already injects.
#[tokio::test]
async fn protocol_mode_is_light_just_like_argv() {
let tmp = TempDir::new().unwrap();
skill_corpus(tmp.path());
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());

let out = inject_first_message_prefix(
"Do stuff",
&mgr,
InjectionConfig {
user_id: "system_default_user",
preset_context: Some("Custom rule"),
skills: &["cron".to_owned()],
delivery_mode: SkillDeliveryMode::Protocol,
},
)
.await;

assert!(out.contains("Custom rule"), "preset context still ships in light mode");
assert!(
!out.contains("Available Skills"),
"protocol delivery must not also inject an index: {out}"
);
}

/// The `injected` block must advertise both channels — this is the wiring
/// half of the truncation-is-safe argument.
#[tokio::test]
async fn injected_mode_carries_the_dual_channel_instructions() {
let tmp = TempDir::new().unwrap();
skill_corpus(tmp.path());
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());

let out = inject_first_message_prefix(
"Hello",
&mgr,
InjectionConfig {
user_id: "system_default_user",
preset_context: None,
skills: &["cron".to_owned()],
delivery_mode: SkillDeliveryMode::Injected,
},
)
.await;
assert!(out.contains("Available Skills"));
assert!(out.contains("skills show"));
assert!(out.contains("[LOAD_SKILL:"));
}

/// `compose_injected_prefix` must produce the block WITHOUT the caller's
/// content, and must not leak the sentinel it splits on.
#[tokio::test]
async fn compose_injected_prefix_returns_the_block_without_content() {
let tmp = TempDir::new().unwrap();
skill_corpus(tmp.path());
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());

let prefix = compose_injected_prefix(
&mgr,
InjectionConfig {
user_id: "system_default_user",
preset_context: Some("Rule 1."),
skills: &["cron".to_owned()],
delivery_mode: SkillDeliveryMode::Injected,
},
)
.await
.expect("a conversation with rules and skills must produce a prefix");

assert!(prefix.contains("[Assistant Rules]"));
assert!(prefix.contains("Rule 1."));
assert!(prefix.contains("Available Skills"));
assert!(
prefix.ends_with("[/Assistant Rules]"),
"trailing blank lines trimmed: {prefix:?}"
);
assert!(!prefix.contains('\u{0}'), "the split sentinel must never escape");
}

/// Nothing to inject must be `None`, not an empty string: the caller uses it
/// to decide whether to touch the prompt at all.
#[tokio::test]
async fn compose_injected_prefix_is_none_when_there_is_nothing_to_inject() {
let tmp = TempDir::new().unwrap();
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());

assert!(
compose_injected_prefix(
&mgr,
InjectionConfig {
user_id: "system_default_user",
preset_context: None,
skills: &[],
delivery_mode: SkillDeliveryMode::Injected,
},
)
.await
.is_none()
);
}

#[tokio::test]
async fn light_mode_with_preset_context() {
async fn argv_mode_is_light_and_only_carries_preset_context() {
let tmp = TempDir::new().unwrap();
let mgr = test_mgr(tmp.path());

Expand All @@ -94,7 +243,7 @@ mod tests {
user_id: "system_default_user",
preset_context: Some("Be concise."),
skills: &[],
native_skill_support: true,
delivery_mode: SkillDeliveryMode::Argv,
},
)
.await;
Expand All @@ -105,7 +254,7 @@ mod tests {
}

#[tokio::test]
async fn light_mode_empty_context_passes_through() {
async fn light_mode_with_no_context_passes_through() {
let tmp = TempDir::new().unwrap();
let mgr = test_mgr(tmp.path());

Expand All @@ -116,15 +265,15 @@ mod tests {
user_id: "system_default_user",
preset_context: None,
skills: &[],
native_skill_support: true,
delivery_mode: SkillDeliveryMode::Argv,
},
)
.await;
assert_eq!(out, "Hello");
}

#[tokio::test]
async fn heavy_mode_no_skills_no_context_passes_through() {
async fn injected_mode_with_no_skills_and_no_context_passes_through() {
let tmp = TempDir::new().unwrap();
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());
Expand All @@ -136,15 +285,15 @@ mod tests {
user_id: "system_default_user",
preset_context: None,
skills: &[],
native_skill_support: false,
delivery_mode: SkillDeliveryMode::Injected,
},
)
.await;
assert_eq!(out, "Hello");
}

#[tokio::test]
async fn heavy_mode_with_preset_context_no_skills() {
async fn injected_mode_with_preset_context_and_no_skills() {
let tmp = TempDir::new().unwrap();
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());
Expand All @@ -156,7 +305,7 @@ mod tests {
user_id: "system_default_user",
preset_context: Some("Rule 1."),
skills: &[],
native_skill_support: false,
delivery_mode: SkillDeliveryMode::Injected,
},
)
.await;
Expand All @@ -167,7 +316,7 @@ mod tests {
}

#[tokio::test]
async fn heavy_mode_with_resolved_skills_injects_index() {
async fn injected_mode_with_resolved_skills_injects_the_index() {
// Set up a builtin skills dir with two skills; pass only one in `skills`.
let tmp = TempDir::new().unwrap();
let auto = tmp.path().join("auto-inject");
Expand All @@ -193,7 +342,7 @@ mod tests {
user_id: "system_default_user",
preset_context: None,
skills: &["cron".to_owned()],
native_skill_support: false,
delivery_mode: SkillDeliveryMode::Injected,
},
)
.await;
Expand All @@ -203,7 +352,7 @@ mod tests {
}

#[tokio::test]
async fn native_support_uses_light_mode_even_with_skills() {
async fn a_layer_one_vendor_stays_light_even_with_skills() {
let tmp = TempDir::new().unwrap();
let _guard = EmptyBuiltinGuard::new(tmp.path());
let mgr = test_mgr(tmp.path());
Expand All @@ -215,7 +364,7 @@ mod tests {
user_id: "system_default_user",
preset_context: Some("Custom rule"),
skills: &["cron".to_owned()],
native_skill_support: true,
delivery_mode: SkillDeliveryMode::Argv,
},
)
.await;
Expand Down
Loading
Loading