use crate::core::ctx::Ctx;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const AGENT_CONTEXT_KEY: &str = "agent_context";
pub const PROJECTION_PLACEMENT_KEY: &str = "__projection_placement";
pub const RUN_ID_KEY: &str = "run_id";
pub const STEP_CTX_KEY: &str = "step_ctx";
pub const PROJECT_NAME_ALIAS_KEY: &str = "project_name_alias";
pub const TASK_PROJECT_ROOT_KEY: &str = "project_root";
pub const TASK_WORK_DIR_KEY: &str = "work_dir";
pub const TASK_METADATA_KEY: &str = "task_metadata";
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
pub struct StepPointer {
pub name: String,
pub size_bytes: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub file_path: Option<String>,
pub content_url: String,
pub sha256: String,
}
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize, schemars::JsonSchema)]
pub struct AgentContextView {
pub task_id: String,
pub agent: String,
pub attempt: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub project_root: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub work_dir: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(with = "Option<serde_json::Value>")]
pub task_metadata: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub run_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub project_name_alias: Option<String>,
#[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
#[schemars(with = "serde_json::Value")]
pub extra: serde_json::Map<String, Value>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub steps: Vec<StepPointer>,
}
impl AgentContextView {
pub fn from_ctx(ctx: &Ctx) -> Self {
let runtime = &ctx.meta.runtime;
Self {
task_id: ctx.task_id.to_string(),
agent: ctx.agent.clone(),
attempt: ctx.attempt,
project_root: runtime
.get(TASK_PROJECT_ROOT_KEY)
.and_then(Value::as_str)
.map(String::from),
work_dir: runtime
.get(TASK_WORK_DIR_KEY)
.and_then(Value::as_str)
.map(String::from),
task_metadata: runtime.get(TASK_METADATA_KEY).cloned(),
run_id: runtime
.get(RUN_ID_KEY)
.and_then(Value::as_str)
.map(String::from),
project_name_alias: runtime
.get(PROJECT_NAME_ALIAS_KEY)
.and_then(Value::as_str)
.map(String::from),
extra: serde_json::Map::new(),
steps: Vec::new(),
}
}
pub fn materialized_or_from_ctx(ctx: &Ctx) -> Self {
ctx.meta
.runtime
.get(AGENT_CONTEXT_KEY)
.and_then(|v| serde_json::from_value::<Self>(v.clone()).ok())
.unwrap_or_else(|| Self::from_ctx(ctx))
}
pub fn to_directive_header(&self) -> String {
let mut out = String::new();
if let Some(alias) = &self.project_name_alias {
out.push_str(&format!("project_name_alias: {alias}\n"));
}
if let Some(root) = &self.project_root {
out.push_str(&format!("project_root: {root}\n"));
}
if let Some(dir) = &self.work_dir {
out.push_str(&format!("work_dir: {dir}\n"));
}
if let Some(meta) = &self.task_metadata {
out.push_str(&format!("task_metadata: {meta}\n"));
}
for (key, value) in &self.extra {
out.push_str(&format!("{key}: {value}\n"));
}
out
}
pub fn apply_policy(mut self, policy: &ContextPolicy) -> Self {
if !policy.allows("project_root") {
self.project_root = None;
}
if !policy.allows("work_dir") {
self.work_dir = None;
}
if !policy.allows("task_metadata") {
self.task_metadata = None;
}
if !policy.allows("run_id") {
self.run_id = None;
}
if !policy.allows("project_name_alias") {
self.project_name_alias = None;
}
self.extra.retain(|key, _| policy.allows(key));
self
}
}
pub use mlua_swarm_schema::ContextPolicy;
#[cfg(test)]
mod tests {
use super::*;
use crate::types::StepId;
fn ctx_with_runtime(pairs: &[(&str, Value)]) -> Ctx {
let mut ctx = Ctx::new(StepId::parse("ST-1").unwrap(), 3, "planner");
for (k, v) in pairs {
ctx.meta.runtime.insert((*k).to_string(), v.clone());
}
ctx
}
#[test]
fn from_ctx_extracts_identity_and_all_runtime_keys() {
let ctx = ctx_with_runtime(&[
(TASK_PROJECT_ROOT_KEY, Value::String("/repo".into())),
(TASK_WORK_DIR_KEY, Value::String("/repo/work".into())),
(TASK_METADATA_KEY, serde_json::json!({"issue": 20})),
(RUN_ID_KEY, Value::String("R-abc".into())),
(PROJECT_NAME_ALIAS_KEY, Value::String("alias-x".into())),
]);
let view = AgentContextView::from_ctx(&ctx);
assert_eq!(view.task_id, "ST-1");
assert_eq!(view.agent, "planner");
assert_eq!(view.attempt, 3);
assert_eq!(view.project_root.as_deref(), Some("/repo"));
assert_eq!(view.work_dir.as_deref(), Some("/repo/work"));
assert_eq!(view.task_metadata, Some(serde_json::json!({"issue": 20})));
assert_eq!(view.run_id.as_deref(), Some("R-abc"));
assert_eq!(view.project_name_alias.as_deref(), Some("alias-x"));
assert!(view.extra.is_empty());
}
#[test]
fn from_ctx_absent_runtime_keys_stay_none() {
let ctx = ctx_with_runtime(&[]);
let view = AgentContextView::from_ctx(&ctx);
assert!(view.project_root.is_none());
assert!(view.work_dir.is_none());
assert!(view.task_metadata.is_none());
assert!(view.run_id.is_none());
assert!(view.project_name_alias.is_none());
}
#[test]
fn materialized_or_from_ctx_prefers_materialized_view() {
let mut view = AgentContextView::from_ctx(&ctx_with_runtime(&[]));
view.task_id = "ST-materialized".into();
view.extra
.insert("custom".into(), Value::String("value".into()));
let mut ctx = ctx_with_runtime(&[]);
ctx.meta.runtime.insert(
AGENT_CONTEXT_KEY.to_string(),
serde_json::to_value(&view).unwrap(),
);
let resolved = AgentContextView::materialized_or_from_ctx(&ctx);
assert_eq!(resolved.task_id, "ST-materialized");
assert_eq!(
resolved.extra.get("custom"),
Some(&Value::String("value".into()))
);
}
#[test]
fn materialized_or_from_ctx_falls_back_when_key_absent() {
let ctx = ctx_with_runtime(&[(TASK_PROJECT_ROOT_KEY, Value::String("/repo".into()))]);
let resolved = AgentContextView::materialized_or_from_ctx(&ctx);
assert_eq!(resolved.project_root.as_deref(), Some("/repo"));
}
#[test]
fn materialized_or_from_ctx_falls_back_when_value_malformed() {
let mut ctx = ctx_with_runtime(&[]);
ctx.meta.runtime.insert(
AGENT_CONTEXT_KEY.to_string(),
Value::String("not an object".into()),
);
let resolved = AgentContextView::materialized_or_from_ctx(&ctx);
assert_eq!(resolved.task_id, "ST-1");
}
#[test]
fn serde_round_trip_skips_absent_optionals_and_empty_extra() {
let view = AgentContextView {
task_id: "ST-1".into(),
agent: "planner".into(),
attempt: 1,
..Default::default()
};
let json = serde_json::to_value(&view).unwrap();
let obj = json.as_object().unwrap();
assert_eq!(
obj.len(),
3,
"only identity fields should serialize: {obj:?}"
);
assert!(obj.contains_key("task_id"));
assert!(obj.contains_key("agent"));
assert!(obj.contains_key("attempt"));
let round_tripped: AgentContextView = serde_json::from_value(json).unwrap();
assert_eq!(round_tripped, view);
}
#[test]
fn serde_round_trip_full_view() {
let mut view = AgentContextView::from_ctx(&ctx_with_runtime(&[
(TASK_PROJECT_ROOT_KEY, Value::String("/repo".into())),
(RUN_ID_KEY, Value::String("R-1".into())),
]));
view.extra
.insert("flow_ir_ctx".into(), serde_json::json!({"k": "v"}));
let json = serde_json::to_value(&view).unwrap();
let round_tripped: AgentContextView = serde_json::from_value(json).unwrap();
assert_eq!(round_tripped, view);
}
#[test]
fn json_schema_export_contains_all_fields() {
let schema = schemars::schema_for!(AgentContextView);
let json = serde_json::to_value(&schema).unwrap();
let props = json["properties"].as_object().expect("properties object");
for field in [
"task_id",
"agent",
"attempt",
"project_root",
"work_dir",
"task_metadata",
"run_id",
"project_name_alias",
"extra",
"steps",
] {
assert!(props.contains_key(field), "schema missing field {field}");
}
}
#[test]
fn context_policy_default_is_pass_all() {
let view = AgentContextView::from_ctx(&ctx_with_runtime(&[
(TASK_PROJECT_ROOT_KEY, Value::String("/repo".into())),
(RUN_ID_KEY, Value::String("R-1".into())),
]));
let filtered = view.clone().apply_policy(&ContextPolicy::default());
assert_eq!(filtered, view);
}
#[test]
fn context_policy_include_only_keeps_listed_fields() {
let view = AgentContextView::from_ctx(&ctx_with_runtime(&[
(TASK_PROJECT_ROOT_KEY, Value::String("/repo".into())),
(TASK_WORK_DIR_KEY, Value::String("/repo/work".into())),
(RUN_ID_KEY, Value::String("R-1".into())),
]));
let policy = ContextPolicy {
include: Some(vec!["project_root".to_string()]),
exclude: Vec::new(),
..Default::default()
};
let filtered = view.apply_policy(&policy);
assert_eq!(filtered.project_root.as_deref(), Some("/repo"));
assert!(filtered.work_dir.is_none());
assert!(filtered.run_id.is_none());
}
#[test]
fn context_policy_exclude_drops_listed_fields() {
let view = AgentContextView::from_ctx(&ctx_with_runtime(&[
(TASK_PROJECT_ROOT_KEY, Value::String("/repo".into())),
(TASK_WORK_DIR_KEY, Value::String("/repo/work".into())),
]));
let policy = ContextPolicy {
include: None,
exclude: vec!["work_dir".to_string()],
..Default::default()
};
let filtered = view.apply_policy(&policy);
assert_eq!(filtered.project_root.as_deref(), Some("/repo"));
assert!(filtered.work_dir.is_none());
}
#[test]
fn context_policy_exclude_beats_include() {
let view = AgentContextView::from_ctx(&ctx_with_runtime(&[(
TASK_PROJECT_ROOT_KEY,
Value::String("/repo".into()),
)]));
let policy = ContextPolicy {
include: Some(vec!["project_root".to_string()]),
exclude: vec!["project_root".to_string()],
..Default::default()
};
let filtered = view.apply_policy(&policy);
assert!(filtered.project_root.is_none());
}
#[test]
fn context_policy_never_filters_identity_fields() {
let view = AgentContextView::from_ctx(&ctx_with_runtime(&[]));
let policy = ContextPolicy {
include: Some(vec![]),
exclude: vec![
"task_id".to_string(),
"agent".to_string(),
"attempt".to_string(),
],
..Default::default()
};
let filtered = view.clone().apply_policy(&policy);
assert_eq!(filtered.task_id, view.task_id);
assert_eq!(filtered.agent, view.agent);
assert_eq!(filtered.attempt, view.attempt);
}
#[test]
fn context_policy_exclude_removes_extra_key() {
let mut view = AgentContextView::from_ctx(&ctx_with_runtime(&[]));
view.extra
.insert("secret".into(), Value::String("x".into()));
view.extra.insert("kept".into(), Value::String("y".into()));
let policy = ContextPolicy {
include: None,
exclude: vec!["secret".to_string()],
..Default::default()
};
let filtered = view.apply_policy(&policy);
assert!(!filtered.extra.contains_key("secret"));
assert!(filtered.extra.contains_key("kept"));
}
#[test]
fn to_directive_header_renders_task_metadata_and_omits_absent_fields() {
let view = AgentContextView {
task_id: "ST-1".into(),
agent: "planner".into(),
attempt: 1,
project_root: Some("/repo".into()),
work_dir: None,
task_metadata: Some(serde_json::json!({"issue": 20})),
run_id: None,
project_name_alias: None,
extra: serde_json::Map::new(),
steps: Vec::new(),
};
let header = view.to_directive_header();
assert_eq!(
header,
"project_root: /repo\ntask_metadata: {\"issue\":20}\n"
);
assert!(!header.contains("work_dir:"));
assert!(!header.contains("project_name_alias:"));
}
#[test]
fn to_directive_header_renders_alias_root_work_dir_in_existing_order() {
let view = AgentContextView {
task_id: "ST-1".into(),
agent: "planner".into(),
attempt: 1,
project_root: Some("/repo".into()),
work_dir: Some("/repo/work".into()),
task_metadata: None,
run_id: None,
project_name_alias: Some("alias-x".into()),
extra: serde_json::Map::new(),
steps: Vec::new(),
};
let header = view.to_directive_header();
assert_eq!(
header,
"project_name_alias: alias-x\nproject_root: /repo\nwork_dir: /repo/work\n"
);
}
#[test]
fn to_directive_header_empty_view_renders_empty_string() {
let view = AgentContextView {
task_id: "ST-1".into(),
agent: "planner".into(),
attempt: 1,
..Default::default()
};
assert_eq!(view.to_directive_header(), "");
}
#[test]
fn extra_field_propagates_to_directive_header_and_serialized_json() {
let mut view = AgentContextView {
task_id: "ST-1".into(),
agent: "planner".into(),
attempt: 1,
..Default::default()
};
view.extra
.insert("step_meta".into(), serde_json::json!({"loop_idx": 2}));
let header = view.to_directive_header();
assert_eq!(header, "step_meta: {\"loop_idx\":2}\n");
let json = serde_json::to_value(&view).unwrap();
assert_eq!(
json.get("extra").and_then(|e| e.get("step_meta")),
Some(&serde_json::json!({"loop_idx": 2}))
);
}
#[test]
fn context_policy_path_compat_reexport() {
let policy: crate::core::agent_context::ContextPolicy =
crate::core::agent_context::ContextPolicy::default();
assert_eq!(policy, mlua_swarm_schema::ContextPolicy::default());
}
}