Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: bind Plugin receipts to durable Host Tool-call history identity

Browse Source
This commit is contained in:
Hare
2026-10-30 09:06:29 +09:00
parent 5329a93e3e
commit 5a317e6d2a
11 changed files with 744 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+76 -51
View File
@@ -1480,15 +1480,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
fn get_unanswered_tool_calls(&self, history: &History<A>) -> (Vec<ToolCall>, Vec<String>) {
let mut pending_calls = Vec::new();
let mut started_call_ids = Vec::new();
let mut answered_call_ids = std::collections::HashSet::new();
for item in history.items() {
if let Item::ToolResult { call_id, .. } = item {
answered_call_ids.insert(call_id.clone());
}
}
for item in history.items() {
for entry in unanswered_tool_call_entries(history) {
if let Item::ToolCall {
call_id,
name,
@@ -1496,8 +1488,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
execution_id,
status,
..
} = item
&& !answered_call_ids.contains(call_id)
} = &entry.item
{
if execution_id.is_some() || *status == Some(ItemStatus::InProgress) {
started_call_ids.push(call_id.clone());
@@ -1639,6 +1630,30 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
} = info;
tool_call = call;
context.call_id = tool_call.id.clone();
if let PreToolAction::SyntheticResult(result) = &action
&& result.tool_use_id != expected_tool_use_id
{
return Err(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
InterceptorError::new(
InterceptorErrorCategory::ContractViolation,
"synthetic tool result changed immutable tool call identity",
),
)
.into());
}
let execution_id = if matches!(action, PreToolAction::Pause) {
None
} else {
Some(context.execution_id())
};
self.append_early_tool_call(history, annotate, &tool_call, call_index, execution_id)?;
// Policy runs before this call exists in history. Only the successful
// append's annotation may bind the identity, overriding hook changes.
context.logical_call_id = history
.entries()
.last()
.and_then(|entry| self.interceptor.logical_tool_call_id(&entry.annotation));
batch.call_info_map.insert(
tool_call.id.clone(),
(tool_call.clone(), meta, tool.clone(), context.clone()),
@@ -1647,13 +1662,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
match action {
PreToolAction::Continue => {
let attempt_id = context.execution_id();
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
Some(attempt_id.clone()),
)?;
batch
.attempt_fence
.register(tool_call.id.clone(), attempt_id.clone());
@@ -1695,34 +1703,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
batch.dispatched_any = true;
}
PreToolAction::SyntheticResult(result) => {
if result.tool_use_id != expected_tool_use_id {
return Err(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
InterceptorError::new(
InterceptorErrorCategory::ContractViolation,
"synthetic tool result changed immutable tool call identity",
),
)
.into());
}
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
Some(context.execution_id()),
)?;
batch.synthetic_results.push_back(result);
batch.dispatched_any = true;
}
PreToolAction::Skip => {
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
Some(context.execution_id()),
)?;
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
"Tool execution skipped by interceptor",
@@ -1730,13 +1714,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
batch.dispatched_any = true;
}
PreToolAction::Abort(reason) => {
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
Some(context.execution_id()),
)?;
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
format!("Tool execution aborted: {reason}"),
@@ -1745,7 +1722,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
return Err(EngineError::Aborted(reason));
}
PreToolAction::Pause => {
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.dispatched_any = true;
return Ok(EarlyToolAdmission::Pause);
}
@@ -2010,10 +1986,24 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count);
self.tool_execution_batch_count += 1;
let mut logical_call_ids: HashMap<String, VecDeque<Option<String>>> = HashMap::new();
for entry in unanswered_tool_call_entries(history) {
if let Item::ToolCall { call_id, .. } = &entry.item {
logical_call_ids
.entry(call_id.clone())
.or_default()
.push_back(self.interceptor.logical_tool_call_id(&entry.annotation));
}
}
// Phase 1: Apply pre_tool_call interceptor (determine skip/abort/synthetic result)
let mut approved_calls = Vec::new();
for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() {
let expected_tool_use_id = tool_call.id.clone();
let logical_call_id = logical_call_ids
.get_mut(&tool_call.id)
.and_then(VecDeque::pop_front)
.flatten();
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
if !tool_call.input.is_object() {
synthetic_results.push(ToolResult::error(
@@ -2079,6 +2069,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
let tool_call = info.call;
let mut context = info.context;
context.call_id = tool_call.id.clone();
context.logical_call_id = logical_call_id;
call_info_map.insert(
tool_call.id.clone(),
(tool_call, info.meta.clone(), info.tool.clone(), context),
@@ -2098,6 +2089,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
tool_call = info.call;
let mut context = info.context;
context.call_id = tool_call.id.clone();
context.logical_call_id = logical_call_id;
call_info_map.insert(
tool_call.id.clone(),
@@ -2111,7 +2103,8 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
approved_calls.push((tool_call, context, Some(info.tool)));
} else {
// Unknown tools go into approved list as-is (will error at execution)
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
let mut context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
context.logical_call_id = logical_call_id;
approved_calls.push((tool_call, context, None));
}
}
@@ -4333,6 +4326,38 @@ fn item_kind(item: &Item) -> &'static str {
}
}
/// Pair each result with one preceding unanswered occurrence, never with a later
/// call that happens to reuse the provider ID. Keep the exact annotated entries
/// so dispatch and restoration resolve the same host-owned identity.
fn unanswered_tool_call_entries<A>(history: &History<A>) -> Vec<&HistoryEntry<A>> {
let mut calls = Vec::new();
let mut by_provider_id: HashMap<&str, VecDeque<usize>> = HashMap::new();
for entry in history.entries() {
match &entry.item {
Item::ToolCall { call_id, .. } => {
by_provider_id
.entry(call_id.as_str())
.or_default()
.push_back(calls.len());
calls.push(Some(entry));
}
Item::ToolResult { call_id, .. } => {
if let Some(queue) = by_provider_id.get_mut(call_id.as_str())
&& let Some(index) = queue.pop_front()
{
calls[index] = None;
}
}
_ => {}
}
}
calls.into_iter().flatten().collect()
}
#[cfg(test)]
#[path = "tests/tool_call_identity.rs"]
mod tool_call_identity_tests;
#[cfg(test)]
mod tests {
use super::*;
crates/agen/src/interceptor.rs
+10 -2
View File
@@ -309,7 +309,7 @@ pub struct ToolCallInfo<'a, A = ()> {
pub meta: ToolMeta,
/// Tool instance (for state access).
pub tool: Arc<dyn Tool>,
/// Response-local execution context for this call.
/// Execution context. The Engine binds durable identity after pre-tool policy.
pub context: ToolExecutionContext,
}
@@ -324,7 +324,7 @@ pub struct ToolResultInfo<'a, A = ()> {
pub meta: ToolMeta,
/// Tool instance (for state access).
pub tool: Arc<dyn Tool>,
/// Response-local execution context for this call.
/// Execution context. The Engine binds durable identity after pre-tool policy.
pub context: ToolExecutionContext,
}
@@ -344,6 +344,14 @@ pub struct ToolResultInfo<'a, A = ()> {
/// flows, permission checks, and other trusted host adaptation.
#[async_trait]
pub trait Interceptor<A: Send + Sync = ()>: Send + Sync {
/// Pure projection of a committed ToolCall's host annotation into an opaque,
/// durable call identity. Agen neither interprets storage metadata nor derives
/// this identity from provider IDs. Early pre-tool policy runs before commit;
/// the Engine binds this value after policy, before execution.
fn logical_tool_call_id(&self, _annotation: &A) -> Option<String> {
None
}
/// Whether side effects may start before [`Self::on_assistant_turn_end`].
///
/// The conservative default preserves the historical whole-response gate.
crates/agen/src/tests/tool_call_identity.rs
+250
View File
@@ -0,0 +1,250 @@
//! Engine component tests with a recording tool and scripted host annotations.
//! Persistence is represented by the append callback; Server/IPC storage is not
//! exercised here. Run with `cargo test -p agen --lib tool_call_identity_tests`.
use super::*;
use async_trait::async_trait;
use std::sync::Mutex;
#[derive(Clone)]
struct NoLlm;
#[async_trait]
impl LlmClient for NoLlm {
async fn stream(&self, _: Request) -> Result<ResponseStream, ClientError> {
panic!("identity tests must not request a model response")
}
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
}
struct HostIdentity {
early: bool,
}
#[async_trait]
impl Interceptor<String> for HostIdentity {
fn logical_tool_call_id(&self, annotation: &String) -> Option<String> {
Some(annotation.clone())
}
async fn pre_tool_call(
&self,
info: &mut ToolCallInfo<'_, String>,
) -> Result<PreToolAction, InterceptorError> {
if self.early {
assert!(info.history.is_empty());
}
// Pre-tool hooks may replace the context, but cannot replace persisted identity.
info.context.logical_call_id = Some("hook-forged-identity".into());
Ok(PreToolAction::Continue)
}
}
struct RecordingTool(Arc<Mutex<Vec<ToolExecutionContext>>>);
#[async_trait]
impl crate::tool::Tool for RecordingTool {
async fn execute(
&self,
_: &str,
context: ToolExecutionContext,
) -> Result<crate::tool::ToolOutput, ToolError> {
self.0.lock().unwrap().push(context);
Ok("recorded".to_string().into())
}
}
fn engine(
early: bool,
) -> (
Engine<NoLlm, Mutable, String>,
Arc<Mutex<Vec<ToolExecutionContext>>>,
) {
let mut engine = Engine::new_annotated(NoLlm);
engine.set_interceptor(HostIdentity { early });
engine.max_turns = Some(0);
engine.start_logical_run();
let observed = Arc::new(Mutex::new(Vec::new()));
let tool = Arc::new(RecordingTool(observed.clone()));
engine.register_tool(Arc::new(move || {
(crate::tool::ToolMeta::new("record"), tool.clone())
}));
engine.tool_server_handle().flush_pending();
(engine, observed)
}
fn entry(item: Item, id: &str) -> HistoryEntry<String> {
HistoryEntry::new(item, id.to_string())
}
fn call(id: &str) -> Item {
Item::tool_call_json(id, "record", json!({}))
}
#[test]
fn results_answer_only_preceding_occurrences_of_reused_provider_ids() {
let history = History::from_entries(vec![
entry(
Item::tool_result("orphan", "no preceding call"),
"orphan-result",
),
entry(call("reused"), "host-first"),
entry(call("sibling"), "host-sibling"),
entry(
Item::tool_result("sibling", "out of order"),
"sibling-result",
),
entry(Item::tool_result("reused", "answered"), "first-result"),
entry(call("reused"), "host-pending"),
entry(call("orphan"), "host-after-orphan"),
]);
let pending = unanswered_tool_call_entries(&history);
assert_eq!(
pending
.iter()
.map(|entry| entry.annotation.as_str())
.collect::<Vec<_>>(),
["host-pending", "host-after-orphan"]
);
}
#[tokio::test]
async fn early_dispatch_binds_identity_only_after_successful_append() {
for reject_append in [false, true] {
let (mut engine, observed) = engine(true);
let mut history = History::new();
let mut batch = EarlyToolExecutionBatch::new("live-batch".into(), 0);
let result = engine
.admit_early_tool_calls(
&mut history,
&mut |_| {
if reject_append {
Err("storage unavailable".into())
} else {
Ok("host-committed-entry".into())
}
},
&mut batch,
vec![crate::timeline::CollectedToolCall {
call_index: 0,
call: ToolCall {
id: "provider-call".into(),
name: "record".into(),
input: json!({}),
},
}],
)
.await;
if reject_append {
assert!(matches!(result, Err(EngineError::HistoryAppend(_))));
assert!(history.is_empty());
assert!(batch.execution_handles.is_empty());
assert!(observed.lock().unwrap().is_empty());
} else {
assert!(matches!(result, Ok(EarlyToolAdmission::Continue)));
let (_, terminal) = batch.futures.next().await.unwrap();
assert_eq!(terminal.tool_use_id, "provider-call");
let observed = observed.lock().unwrap();
assert_eq!(observed.len(), 1);
assert_eq!(
observed[0].logical_call_id.as_deref(),
Some(history.entries()[0].annotation.as_str())
);
assert_eq!(
observed[0].logical_call_id.as_deref(),
Some("host-committed-entry")
);
assert_eq!(batch.call_info_map["provider-call"].3, observed[0]);
}
}
}
#[tokio::test]
async fn normal_dispatch_and_unstarted_restore_use_the_exact_unanswered_annotation() {
for restore in [false, true] {
let (mut engine, observed) = engine(false);
let mut history = History::from_entries(vec![
entry(call("reused"), "host-answered"),
entry(
Item::tool_result("reused", "already completed"),
"answered-result",
),
]);
engine
.append_history_items(&mut history, [call("reused")], &mut |_| {
Ok("host-pending".into())
})
.unwrap();
if restore {
let persisted = serde_json::to_string(&history).unwrap();
history = serde_json::from_str(&persisted).unwrap();
}
let (pending, started) = engine.get_unanswered_tool_calls(&history);
assert_eq!(pending.len(), 1);
assert!(started.is_empty());
if restore {
engine
.run_turn_loop(&mut history, &mut |_| Ok("terminal-result".into()))
.await
.unwrap();
} else {
let result = engine
.execute_tools(&mut history, &mut |_| Ok("terminal-result".into()), pending)
.await
.unwrap();
assert!(matches!(result, ToolExecutionResult::Completed));
}
assert_eq!(observed.lock().unwrap().len(), 1);
let context = observed.lock().unwrap()[0].clone();
assert_eq!(context.call_id, "reused");
assert_eq!(context.logical_call_id.as_deref(), Some("host-pending"));
// A subsequent resume sees the terminal result, not another live attempt.
engine
.run_turn_loop(&mut history, &mut |_| Ok("unexpected-append".into()))
.await
.unwrap();
assert_eq!(observed.lock().unwrap().len(), 1);
assert!(unanswered_tool_call_entries(&history).is_empty());
}
}
#[tokio::test]
async fn started_restore_with_reused_provider_id_closes_unknown_without_reexecution() {
let (mut engine, observed) = engine(false);
let mut history = History::from_entries(vec![
entry(call("reused"), "host-answered"),
entry(
Item::tool_result("reused", "already completed"),
"answered-result",
),
entry(
call("reused")
.with_tool_execution_metadata(0, Some("old-attempt".into()))
.with_status(ItemStatus::InProgress),
"host-started",
),
]);
let (pending, started) = engine.get_unanswered_tool_calls(&history);
assert!(pending.is_empty());
assert_eq!(started, ["reused"]);
engine
.run_turn_loop(&mut history, &mut |_| Ok("unknown-result".into()))
.await
.unwrap();
assert!(observed.lock().unwrap().is_empty());
assert!(matches!(
&history.entries().last().unwrap().item,
Item::ToolResult { call_id, disposition: ToolResultDisposition::OutcomeUnknown, .. }
if call_id == "reused"
));
let len = history.len();
engine
.run_turn_loop(&mut history, &mut |_| Ok("unexpected-append".into()))
.await
.unwrap();
assert_eq!(history.len(), len);
assert!(observed.lock().unwrap().is_empty());
}
crates/agen/src/tool.rs
+14 -6
View File
@@ -336,15 +336,16 @@ pub type ToolDefinition = Arc<dyn Fn() -> (ToolMeta, Arc<dyn Tool>) + Send + Syn
/// Per-call context supplied by the engine when executing a tool call.
///
/// The context identifies a tool call within one assistant response's tool-call
/// batch without imposing any scheduling policy on the engine. Tool
/// implementations may use it for response-local ordering, diagnostics, or
/// correlation, but it is intentionally not a handle to engine state, history,
/// or session mutation.
/// Live batch identity and ordering are response-local. An optional opaque host
/// identity comes from the committed ToolCall annotation and survives restore.
/// Neither is a handle to engine state, history, or session mutation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolExecutionContext {
/// Provider/tool-call id for the call being executed.
pub call_id: String,
/// Opaque host identity from the committed ToolCall annotation, stable across
/// restoration. Provider call IDs and live batch IDs are not replay authority.
pub logical_call_id: Option<String>,
/// Engine-local identity shared by all tool calls from one execution batch.
pub batch_id: String,
/// Zero-based order of this call in the model-returned tool-call list.
@@ -355,11 +356,18 @@ impl ToolExecutionContext {
pub fn new(call_id: impl Into<String>, batch_id: impl Into<String>, call_index: usize) -> Self {
Self {
call_id: call_id.into(),
logical_call_id: None,
batch_id: batch_id.into(),
call_index,
}
}
/// Bind an identity supplied by the host's persisted-call annotation.
pub fn with_logical_call_id(mut self, logical_call_id: impl Into<String>) -> Self {
self.logical_call_id = Some(logical_call_id.into());
self
}
/// Identifies one live execution attempt without making the batch id a durable
/// replay or idempotency authority.
pub fn execution_id(&self) -> String {
@@ -572,7 +580,7 @@ pub trait Tool: Send + Sync {
///
/// # Arguments
/// * `input_json` - JSON-formatted arguments generated by LLM
/// * `ctx` - response-local call identity and ordering context
/// * `ctx` - per-call identity and live execution ordering context
///
/// # Returns
/// A [`ToolOutput`] with summary and optional detailed content.
crates/worker/src/interrupt_prep.rs
+28 -9
View File
@@ -18,18 +18,25 @@ use agen::{Item, ToolResultDisposition};
/// Build synthetic `Item::ToolResult` items for every unanswered
/// `Item::ToolCall` in `history`, preserving order.
pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Vec<Item> {
let mut answered: std::collections::HashSet<&str> = std::collections::HashSet::new();
let mut pending: Vec<Option<&str>> = Vec::new();
for item in history {
if let Item::ToolResult { call_id, .. } = item {
answered.insert(call_id.as_str());
match item {
Item::ToolCall { call_id, .. } => pending.push(Some(call_id)),
Item::ToolResult { call_id, .. } => {
if let Some(index) = pending
.iter()
.position(|entry| *entry == Some(call_id.as_str()))
{
pending[index] = None;
}
}
_ => {}
}
}
let mut out = Vec::new();
for item in history {
if let Item::ToolCall { call_id, .. } = item {
if !answered.contains(call_id.as_str()) {
for call_id in pending.into_iter().flatten() {
out.push(Item::tool_result_item_with_disposition_and_attachments(
call_id.clone(),
call_id,
summary,
Some(
"Execution ended before completion could be confirmed. Completion and side effects are unknown."
@@ -39,8 +46,6 @@ pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Ve
Vec::new(),
));
}
}
}
out
}
@@ -59,6 +64,20 @@ fn interrupt_tool_result_summary() -> String {
mod tests {
use super::*;
#[test]
fn answered_provider_id_does_not_hide_later_unanswered_logical_call() {
let history = vec![
Item::tool_call("reused", "plugin", "{}"),
Item::tool_result("reused", "first finished"),
Item::tool_call("reused", "plugin", "{}"),
];
let closed = orphan_tool_result_closures(&history, "Interrupted");
assert_eq!(closed.len(), 1);
assert!(
matches!(&closed[0], Item::ToolResult { call_id, disposition: ToolResultDisposition::OutcomeUnknown, .. } if call_id == "reused")
);
}
#[test]
fn no_orphans_returns_empty() {
let history = vec![Item::user_message("hi"), Item::assistant_message("hello")];
crates/worker/src/ipc/interceptor.rs
+4
View File
@@ -437,6 +437,10 @@ impl WorkerInterceptor {
#[async_trait]
impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
fn logical_tool_call_id(&self, annotation: &SessionHistoryMetadata) -> Option<String> {
Some(annotation.entry_id.0.clone())
}
fn supports_early_tool_dispatch(&self) -> bool {
self.registry.on_turn_end.is_empty()
}
crates/worker/src/plugin.rs
+8 -6
View File
@@ -291,18 +291,20 @@ impl Tool for PluginTool {
input_json: &str,
context: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
if context.call_id.is_empty() || context.call_id == "direct" {
return Err(failure("Plugin Tool requires committed Tool-call identity"));
}
let logical_call_id = context
.logical_call_id
.as_deref()
.filter(|id| !id.is_empty())
.ok_or_else(|| failure("Plugin Tool requires committed logical Tool-call identity"))?;
let _capacity = self
.capacity
.clone()
.try_acquire_owned()
.map_err(|_| failure("Plugin invocation capacity exceeded"))?;
let lease = self.start(&context)?;
// The persisted Session and provider call ID survive Runtime restoration;
// the engine-local batch ID identifies only cancellation of a live attempt.
let id_bytes = serde_json::to_vec(&(&self.session_id, &context.call_id))
// Only the host-owned persisted call identity is durable. Provider IDs may
// repeat in a Session; batch/provider IDs identify live attempts only.
let id_bytes = serde_json::to_vec(&(&self.session_id, logical_call_id))
.expect("identity serialization");
let hex: String = Sha256::digest(&id_bytes)
.iter()
crates/worker/src/plugin/tests.rs
+66 -7
View File
@@ -320,6 +320,7 @@ fn install(feature: PluginToolsFeature) -> (Engine<NoLlm, Mutable, ()>, HookRegi
fn context(call: &str, batch: &str) -> ToolExecutionContext {
ToolExecutionContext::new(call, batch, 0)
.with_logical_call_id(format!("persisted-entry-{call}"))
}
fn stopping() -> WorkerStoppingContext {
@@ -502,15 +503,41 @@ async fn adapter_late_publish_denial_withholds_successful_guest_output() {
}
#[tokio::test]
async fn consumer_identity_survives_batch_changes_but_isolates_sessions_and_calls() {
async fn consumer_identity_uses_host_entry_not_provider_call_or_live_batch() {
let fixture = Fixture::new();
let owner = Owner::new(Behavior::Allow);
// Re-preparation simulates a restored adapter, not just reuse of a live Tool.
for (session, call, batch) in [
("session-a", "provider-call", "batch-before-restore"),
("session-a", "provider-call", "batch-after-restore"),
("session-b", "provider-call", "batch-after-restore"),
("session-a", "different-call", "batch-after-restore"),
for (session, call, logical, batch) in [
(
"session-a",
"provider-call",
"host-entry-a",
"batch-before-restore",
),
(
"session-a",
"provider-call",
"host-entry-a",
"batch-after-restore",
),
(
"session-b",
"provider-call",
"host-entry-a",
"batch-after-restore",
),
(
"session-a",
"provider-call",
"host-entry-b",
"batch-after-restore",
),
(
"session-a",
"different-provider-call",
"host-entry-a",
"different-batch",
),
] {
let feature = fixture
.prepare(&["installation_a"], owner.clone(), session)
@@ -518,7 +545,11 @@ async fn consumer_identity_survives_batch_changes_but_isolates_sessions_and_call
let (engine, _) = install(feature);
engine
.tool_server_handle()
.call_tool(TOOL_A, INPUT, context(call, batch))
.call_tool(
TOOL_A,
INPUT,
ToolExecutionContext::new(call, batch, 0).with_logical_call_id(logical),
)
.await
.unwrap();
}
@@ -526,6 +557,34 @@ async fn consumer_identity_survives_batch_changes_but_isolates_sessions_and_call
assert_eq!(consumers[0], consumers[1]);
assert_ne!(consumers[0], consumers[2]);
assert_ne!(consumers[0], consumers[3]);
assert_eq!(consumers[0], consumers[4]);
}
#[tokio::test]
async fn missing_or_empty_logical_identity_refuses_owner_admission() {
let fixture = Fixture::new();
let owner = Owner::new(Behavior::Allow);
let feature = fixture
.prepare(&["installation_a"], owner.clone(), "session-a")
.await;
let (engine, _) = install(feature);
let tools = engine.tool_server_handle();
for context in [
ToolExecutionContext::direct(),
ToolExecutionContext::new("provider-call", "batch", 0),
ToolExecutionContext::new("provider-call", "batch", 0).with_logical_call_id(""),
] {
let error = tools.call_tool(TOOL_A, INPUT, context).await.unwrap_err();
assert!(
error
.to_string()
.contains("committed logical Tool-call identity")
);
}
assert!(
owner.requests().is_empty(),
"identity rejection must precede Begin"
);
}
struct FailInstall;
crates/workspace-server/src/server_plugin_controller_tests.rs
+145 -15
View File
@@ -13,13 +13,14 @@
// Run: cargo test -p yoi-workspace-server --lib plugin_controller_tests
// Prerequisites: openssl and the normal connected-restore test environment.
mod plugin_controller_tests {
use super::plugin_host_tests::plugin_runtime_tests::{Model, assert_https_history};
use super::plugin_host_tests::plugin_runtime_tests::{
Model, assert_https_history, persisted_call_consumers,
};
use super::plugin_tests::request;
use super::*;
use crate::plugin_https::test_support::{HOST, Peer, response};
use crate::plugin_plan::PlanPinState;
use session_store::Store;
use sha2::Digest;
use std::sync::{
Mutex,
atomic::{AtomicUsize, Ordering},
@@ -41,6 +42,10 @@ mod plugin_controller_tests {
impl Fixture {
async fn new(peer: &Peer) -> Self {
Self::with_request_count(peer, 1).await
}
async fn with_request_count(peer: &Peer, request_count: usize) -> Self {
// Use the established short-path/environment guard for real restore
// (not a mock controller or a replacement execution backend).
let directory = super::restored_parent_control_tests::connected_restore_tempdir();
@@ -57,7 +62,7 @@ mod plugin_controller_tests {
let clients = model_clients.clone();
let observed = requests.clone();
let tool = model_tool.clone();
let transport = Arc::new(peer.transport());
let transport = Arc::new(peer.transport_for_requests(request_count));
let api = WorkspaceApi::new_with_embedded_configurators(
config,
store,
@@ -345,6 +350,14 @@ mod plugin_controller_tests {
&self,
worker: &RuntimeWorkerRef,
) -> (session_store::SessionId, Vec<agen::Item>) {
let (session, saved) = self.saved_state(worker);
(session, saved.history)
}
fn saved_state(
&self,
worker: &RuntimeWorkerRef,
) -> (session_store::SessionId, session_store::RestoredState) {
let reopened =
session_store::WorkerSessionStore::new(self.aggregate(worker).join("session"))
.unwrap();
@@ -352,7 +365,7 @@ mod plugin_controller_tests {
let segments = reopened.list_segments(session).unwrap();
assert_eq!(segments.len(), 1);
let saved = session_store::restore(&reopened, session, segments[0]).unwrap();
(session, saved.history)
(session, saved)
}
fn count(&self, table: &str) -> i64 {
@@ -444,10 +457,14 @@ mod plugin_controller_tests {
}
async fn submit(connection: &mut WorkspaceWorkerProtocolConnection) {
submit_id(connection, "controller-https-submit").await;
}
async fn submit_id(connection: &mut WorkspaceWorkerProtocolConnection, request_id: &str) {
connection
.methods
.send(protocol::Method::Submit {
submission_request_id: "controller-https-submit".into(),
submission_request_id: request_id.into(),
input: vec![protocol::Segment::Text {
content: "Check the service status.".into(),
}],
@@ -464,6 +481,14 @@ mod plugin_controller_tests {
async fn finish(
connection: &mut WorkspaceWorkerProtocolConnection,
expected: protocol::RunResult,
) {
finish_id(connection, expected, "controller-https-submit").await;
}
async fn finish_id(
connection: &mut WorkspaceWorkerProtocolConnection,
expected: protocol::RunResult,
request_id: &str,
) {
tokio::time::timeout(WATCHDOG, async {
let mut accepted = false;
@@ -474,7 +499,7 @@ mod plugin_controller_tests {
submission_request_id,
..
} => {
assert_eq!(submission_request_id, "controller-https-submit");
assert_eq!(submission_request_id, request_id);
accepted = true;
}
protocol::Event::SubmissionRejected { message, .. } => {
@@ -489,7 +514,7 @@ mod plugin_controller_tests {
if finished
&& snapshot.state == protocol::WorkerState::Idle
&& snapshot.last_finished_submission_request_id.as_deref()
== Some("controller-https-submit") =>
== Some(request_id) =>
{
return;
}
@@ -501,6 +526,116 @@ mod plugin_controller_tests {
.expect("real Controller HTTPS submission did not finish");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn repeated_provider_call_id_has_distinct_durable_receipts_and_completed_restore_never_reruns()
{
// Gemini also reuses a response-local id for the same Tool. Deliberately
// emit the identical provider id in two explicit user submissions.
let peer = Peer::new(
HOST,
vec![response(503, &[], b"old"), response(200, &[], b"fresh")],
)
.await;
let f = Fixture::with_request_count(&peer, 2).await;
f.install(&peer).await;
let worker = f.create(PROFILE).await;
let public_name = f.model_tool.lock().unwrap().clone();
let mut connection = connect_workspace_worker_protocol(&f.api, &worker, None)
.await
.unwrap();
snapshot(&mut connection).await;
for request in ["first-explicit-call", "second-explicit-call"] {
submit_id(&mut connection, request).await;
finish_id(&mut connection, protocol::RunResult::Finished, request).await;
}
assert_eq!(
peer.connection_count(),
2,
"a new logical call must not replay the first receipt; model requests: {:?}",
f.requests.lock().unwrap().iter().map(|r| &r.items).collect::<Vec<_>>()
);
assert_eq!(f.count("plugin_invocations"), 2);
assert_eq!(f.count("plugin_effect_receipts"), 2);
let statuses = |history: &[agen::Item]| -> Vec<u16> {
history
.iter()
.filter_map(|item| match item {
agen::Item::ToolResult {
call_id,
is_error: false,
content: Some(content),
..
} if call_id == "https-call" => Some(
serde_json::from_str::<Value>(content).unwrap()["structured"]["status"]
.as_u64()
.unwrap() as u16,
),
_ => None,
})
.collect()
};
{
let requests = f.requests.lock().unwrap();
assert_eq!(requests.len(), 4);
assert_eq!(statuses(&requests[3].items), [503, 200]);
}
drop(connection);
f.stop(&worker).await;
let (session, saved) = f.saved_state(&worker);
assert_eq!(statuses(&saved.history), [503, 200]);
let calls: Vec<_> = saved
.annotated_history
.iter()
.filter_map(|entry| match &entry.item {
session_store::LoggedItem::ToolCall { call_id, name, .. }
if name == &public_name =>
{
Some((call_id, &entry.metadata.entry_id))
}
_ => None,
})
.collect();
assert_eq!(calls.len(), 2);
assert_eq!(
calls[0].0, calls[1].0,
"provider id must actually be reused in the fixture"
);
assert_ne!(calls[0].1, calls[1].1);
let mut expected = persisted_call_consumers(session, &saved, &public_name);
expected.sort();
let consumers: Vec<String> = f.api.config_store.with_conn(|conn| {
let mut statement = conn.prepare("SELECT json_extract(consumer_json,'$.tool_call_id') FROM plugin_invocations WHERE workspace_id=?1 AND worker_id=?2 ORDER BY 1")?;
Ok(statement.query_map(rusqlite::params![TEST_WORKSPACE_ID, worker.worker_id], |row| row.get(0))?.collect::<std::result::Result<Vec<_>,_>>()?)
}).unwrap();
assert_eq!(
consumers, expected,
"receipts must bind committed Host annotation ids, not response-local ids"
);
let restored = f.restore(&worker, "completed-reused-id-restore").await;
assert_eq!(
restored.state,
server_api::WorkerRestoreState::Accepted,
"{restored:?}"
);
assert_eq!(f.requests.lock().unwrap().len(), 4);
let mut connection = connect_workspace_worker_protocol(&f.api, &worker, None)
.await
.unwrap();
assert!(snapshot(&mut connection).await.tools.contains(&public_name));
assert_eq!(peer.connection_count(), 2);
assert_eq!(f.count("plugin_invocations"), 2);
assert_eq!(f.count("plugin_effect_receipts"), 2);
drop(connection);
f.stop(&worker).await;
let (restored_session, saved) = f.saved_state(&worker);
assert_eq!(restored_session, session);
let mut restored = persisted_call_consumers(session, &saved, &public_name);
restored.sort();
assert_eq!(restored, expected);
assert_eq!(statuses(&saved.history), [503, 200]);
f.close().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn owner_diagnostics_observes_saved_pin_and_live_refusal_without_execution_or_paths() {
let peer = Peer::new(HOST, vec![]).await;
@@ -735,14 +870,9 @@ mod plugin_controller_tests {
rusqlite::params![TEST_WORKSPACE_ID, saved.plan_id, worker.worker_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)))?)
}).unwrap();
let identity = serde_json::to_vec(&(session.to_string(), "https-call")).unwrap();
let consumer_id = format!(
"tool-{}",
sha2::Sha256::digest(identity)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>()
);
let consumers = persisted_call_consumers(session, &f.saved_state(&worker).1, &public_name);
assert_eq!(consumers.len(), 1);
let consumer_id = consumers[0].clone();
assert_eq!(
serde_json::from_str::<server_api::PluginConsumer>(&consumer).unwrap(),
server_api::PluginConsumer::ToolCall {
crates/workspace-server/src/server_plugin_remote_controller_tests.rs
+6 -10
View File
@@ -7,7 +7,9 @@
// Run: cargo test -p yoi-workspace-server --lib plugin_remote_controller_tests
// Prerequisites: openssl and the existing connected Runtime test environment.
mod plugin_remote_controller_tests {
use super::plugin_host_tests::plugin_runtime_tests::{Model, assert_https_history};
use super::plugin_host_tests::plugin_runtime_tests::{
Model, assert_https_history, persisted_call_consumers,
};
use super::plugin_tests::request;
use super::subject_authorization_tests::{CONNECTED_RUNTIME, ConnectedWorkerHarness};
use super::*;
@@ -15,7 +17,6 @@ mod plugin_remote_controller_tests {
use crate::plugin_plan::PlanPinState;
use agen::llm_client::{ClientError, LlmClient, Request as ModelRequest};
use session_store::Store;
use sha2::Digest;
use std::sync::Mutex;
use std::time::Duration;
@@ -380,14 +381,9 @@ mod plugin_remote_controller_tests {
rusqlite::params![TEST_WORKSPACE_ID, saved.plan_id, worker.worker_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)))?)
}).unwrap();
let identity = serde_json::to_vec(&(session.to_string(), "https-call")).unwrap();
let consumer_id = format!(
"tool-{}",
sha2::Sha256::digest(identity)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>()
);
let consumers = persisted_call_consumers(session, &history, &public_name);
assert_eq!(consumers.len(), 1);
let consumer_id = consumers[0].clone();
assert_eq!(
serde_json::from_str::<server_api::PluginConsumer>(&consumer).unwrap(),
server_api::PluginConsumer::ToolCall {
crates/workspace-server/src/server_plugin_runtime_tests.rs
+135 -19
View File
@@ -147,7 +147,7 @@ pub(super) mod plugin_runtime_tests {
turn
};
let events = match turn {
0 => vec![
0 | 2 => vec![
ModelEvent::tool_use_start(0, "https-call", &self.tool_name),
ModelEvent::tool_input_delta(
0,
@@ -159,7 +159,7 @@ pub(super) mod plugin_runtime_tests {
),
ModelEvent::tool_use_stop(0),
],
1 => vec![
1 | 3 => vec![
ModelEvent::text_block_start(0),
ModelEvent::text_delta(0, "Observed the HTTPS tool result."),
ModelEvent::text_block_stop(0, None),
@@ -177,6 +177,32 @@ pub(super) mod plugin_runtime_tests {
}
}
pub(in crate::server::tests) fn persisted_call_consumers(
session: session_store::SessionId,
saved: &session_store::RestoredState,
public_name: &str,
) -> Vec<String> {
saved
.annotated_history
.iter()
.filter_map(|entry| match &entry.item {
session_store::LoggedItem::ToolCall { name, .. } if name == public_name => {
let bytes =
serde_json::to_vec(&(session.to_string(), &entry.metadata.entry_id.0))
.unwrap();
Some(format!(
"tool-{}",
Sha256::digest(bytes)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>()
))
}
_ => None,
})
.collect()
}
pub(in crate::server::tests) fn assert_https_history(items: &[agen::Item], public_name: &str) {
let calls = items
.iter()
@@ -234,7 +260,19 @@ pub(super) mod plugin_runtime_tests {
assert_eq!(output["structured"], json!({"status":503,"body_bytes":9}));
}
async fn https_through_runtime_and_history(remote: bool) {
struct PauseBeforeFirstCall;
#[async_trait::async_trait]
impl worker::hook::Hook<worker::hook::PreToolCall> for PauseBeforeFirstCall {
async fn call(
&self,
_: &worker::hook::ToolCallSummary,
) -> std::result::Result<worker::hook::HookPreToolAction, worker::hook::HookError> {
Ok(worker::hook::HookPreToolAction::Pause)
}
}
async fn https_through_runtime_and_history(remote: bool, restore_pending: bool) {
let mut peer = Peer::new(HOST, vec![response(503, &[], b"peer data")]).await;
let mut f = Fixture::new(Some(peer_binding(&peer))).await;
f.api.plugin_https_transport = Some(Arc::new(peer.transport()));
@@ -263,17 +301,34 @@ pub(super) mod plugin_runtime_tests {
tool_name: public_name.clone(),
requests: Arc::default(),
};
let manifest = worker::WorkerManifest::from_toml(&format!(
let mut manifest = worker::WorkerManifest::from_toml(&format!(
"[worker]\nname=\"{}\"\n[model]\nscheme=\"anthropic\"\nmodel_id=\"scripted\"\n[engine]\n[scope]\nallow=[]\n",
f.source.worker_id
))
.unwrap();
if restore_pending {
manifest
.plugins
.selections
.push(manifest::PluginProfileSelection {
installation_id: f.installed.installation.installation_id.clone(),
tools: vec![f.tool.into()],
logging: false,
https_slots: vec!["api".into()],
config: Default::default(),
});
}
let store_root = f._directory.path().join("tool-history");
let mut host = worker::Worker::new(
manifest,
agen::Engine::new_annotated(model.clone())
.system_prompt("Use the selected HTTPS tool."),
let metadata_root = f._directory.path().join("worker-metadata");
let store = session_store::CombinedStore::new(
session_store::FsStore::new(&store_root).unwrap(),
session_store::FsWorkerStore::new(&metadata_root).unwrap(),
);
let mut host = worker::Worker::new(
manifest.clone(),
agen::Engine::new_annotated(Box::new(model.clone()) as Box<dyn LlmClient>)
.system_prompt("Use the selected HTTPS tool."),
store,
worker::WorkerWorkspaceContext::unavailable(
Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()),
"unrelated Workspace operations are not bound by this fixture",
@@ -330,6 +385,10 @@ pub(super) mod plugin_runtime_tests {
tools.get_tool(f.tool).is_none(),
"bare guest names are not public tools"
);
if restore_pending {
host.enable_worker_metadata_write_through().unwrap();
host.add_pre_tool_call_hook(PauseBeforeFirstCall);
}
host.wire_history_persistence();
let result = tokio::time::timeout(
std::time::Duration::from_secs(30),
@@ -340,7 +399,63 @@ pub(super) mod plugin_runtime_tests {
.await
.expect("HTTPS Tool/Engine must finish within the fixture deadline")
.unwrap();
if restore_pending {
assert!(matches!(result, worker::WorkerRunResult::Paused));
assert_eq!(peer.connection_count(), 0);
assert_eq!(f.count("plugin_invocations"), 0);
let reopened = session_store::FsStore::new(&store_root).unwrap();
let segments = reopened.list_segments(session_id).unwrap();
let pending = session_store::restore(&reopened, session_id, segments[0]).unwrap();
let expected = persisted_call_consumers(session_id, &pending, &public_name);
assert_eq!(
expected.len(),
1,
"paused call must already have a durable Host identity"
);
assert!(
!pending
.history
.iter()
.any(|item| matches!(item, agen::Item::ToolResult { .. }))
);
drop(host);
let mut restored = worker::Worker::prepare_restored_from_worker_metadata_with_context_and_model_client(
&manifest.worker.name, manifest.clone(), session_store::CombinedStore::new(
reopened, session_store::FsWorkerStore::new(&metadata_root).unwrap(),
),
worker::PromptCatalogSource::builtins_only(),
worker::WorkerWorkspaceContext::unavailable(Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()), "unrelated Workspace operations unavailable"),
worker::WorkerFilesystemAuthority::None, Some(Box::new(model.clone())),
).await.unwrap().activate().await.unwrap();
assert_eq!(restored.session_id(), session_id);
let reloaded = session_store::restore(
&session_store::FsStore::new(&store_root).unwrap(),
session_id,
segments[0],
)
.unwrap();
assert_eq!(
persisted_call_consumers(session_id, &reloaded, &public_name),
expected
);
let feature = bridge
.prepare(&worker_ref, &snapshot, &session_id.to_string())
.await
.unwrap();
let report =
restored.install_features(FeatureRegistryBuilder::new().with_module(feature));
assert!(!report.has_errors(), "{}", report.error_message());
restored.wire_history_persistence();
let result =
tokio::time::timeout(std::time::Duration::from_secs(30), restored.resume())
.await
.unwrap()
.unwrap();
assert!(matches!(result, worker::WorkerRunResult::Finished));
host = restored;
} else {
assert!(matches!(result, worker::WorkerRunResult::Finished));
}
assert_https_history(&host.history(), &public_name);
{
let requests = model.requests.lock().unwrap();
@@ -371,14 +486,9 @@ pub(super) mod plugin_runtime_tests {
assert_eq!(segments.len(), 1);
let saved = session_store::restore(&reopened, session_id, segments[0]).unwrap();
assert_https_history(&saved.history, &public_name);
let identity_bytes = serde_json::to_vec(&(session_id.to_string(), "https-call")).unwrap();
let consumer_id = format!(
"tool-{}",
Sha256::digest(identity_bytes)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>()
);
let consumers = persisted_call_consumers(session_id, &saved, &public_name);
assert_eq!(consumers.len(), 1);
let consumer_id = consumers[0].clone();
let (consumer, state, status, body): (String, String, i64, Vec<u8>) = f
.api
.config_store
@@ -418,7 +528,8 @@ pub(super) mod plugin_runtime_tests {
.call_tool(
&public_name,
"{}",
ToolExecutionContext::new("denied-call", "denied-batch", 0),
ToolExecutionContext::new("denied-call", "denied-batch", 0)
.with_logical_call_id("denied-persisted-call"),
)
.await
.is_err()
@@ -444,12 +555,17 @@ pub(super) mod plugin_runtime_tests {
#[tokio::test]
async fn remote_pinned_https_tool_commits_normal_history_and_receipt_then_fences_live_denial() {
https_through_runtime_and_history(true).await;
https_through_runtime_and_history(true, false).await;
}
#[tokio::test]
async fn embedded_pinned_https_tool_commits_normal_history_and_receipt_then_fences_live_denial()
{
https_through_runtime_and_history(false).await;
https_through_runtime_and_history(false, false).await;
}
#[tokio::test]
async fn pending_tool_call_restores_its_committed_host_identity_before_first_effect() {
https_through_runtime_and_history(false, true).await;
}
}