feat: integrate worker submit queue
This commit is contained in:
@@ -120,12 +120,15 @@ mod tests {
|
||||
let mut client = Client::new(socket);
|
||||
|
||||
client
|
||||
.send(&Method::run_text("hello"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"hello",
|
||||
))
|
||||
.await
|
||||
.expect("send method");
|
||||
assert!(matches!(
|
||||
decode_method(&client.socket.sent[0]),
|
||||
Ok(Method::Run { .. })
|
||||
Ok(Method::Submit { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
client.next_event().await,
|
||||
|
||||
@@ -89,12 +89,15 @@ mod tests {
|
||||
let mut client = Client::new(socket);
|
||||
|
||||
client
|
||||
.send(&Method::run_text("hello"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"hello",
|
||||
))
|
||||
.await
|
||||
.expect("send method");
|
||||
assert!(matches!(
|
||||
peer.next().await.as_deref().map(decode_method),
|
||||
Some(Ok(Method::Run { .. }))
|
||||
Some(Ok(Method::Submit { .. }))
|
||||
));
|
||||
|
||||
peer.send(
|
||||
|
||||
@@ -147,12 +147,18 @@ mod tests {
|
||||
|
||||
let mut client = Client::new(Socket::connect(&socket_path).await.unwrap());
|
||||
client
|
||||
.send(&Method::run_text("hello"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"hello",
|
||||
))
|
||||
.await
|
||||
.expect("send method");
|
||||
|
||||
let received = server.await.unwrap().expect("method message");
|
||||
assert!(matches!(decode_method(&received), Ok(Method::Run { .. })));
|
||||
assert!(matches!(
|
||||
decode_method(&received),
|
||||
Ok(Method::Submit { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -114,7 +114,7 @@ mod tests {
|
||||
assert!(matches!(
|
||||
message,
|
||||
Message::Text(ref text)
|
||||
if matches!(decode_method(text), Ok(Method::Run { .. }))
|
||||
if matches!(decode_method(text), Ok(Method::Submit { .. }))
|
||||
));
|
||||
let event = encode_event(&Event::Status {
|
||||
status: WorkerStatus::Idle,
|
||||
@@ -126,7 +126,10 @@ mod tests {
|
||||
let request = format!("ws://{address}").into_client_request().unwrap();
|
||||
let mut client = Client::new(Socket::connect(request).await.unwrap());
|
||||
client
|
||||
.send(&Method::run_text("hello"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"hello",
|
||||
))
|
||||
.await
|
||||
.expect("send method");
|
||||
assert!(matches!(
|
||||
|
||||
+211
-52
@@ -11,6 +11,11 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
pub use identity::{WorkerId, WorkerIdParseError};
|
||||
|
||||
/// Allocate an opaque idempotency key for one client Submit request.
|
||||
pub fn new_submission_request_id() -> String {
|
||||
uuid::Uuid::now_v7().to_string()
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
@@ -27,21 +32,80 @@ fn is_false(value: &bool) -> bool {
|
||||
// Method (Client → Worker via Unix Socket)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Trusted Server → Runtime transport header carrying the authenticated
|
||||
/// browser Account identity for one Worker protocol connection.
|
||||
///
|
||||
/// Runtime accepts this only after its normal HTTP authentication succeeds;
|
||||
/// serialized [`Method`] payloads cannot set authenticated source identity.
|
||||
pub const AUTHENTICATED_ACCOUNT_ID_HEADER: &str = "x-yoi-authenticated-account-id";
|
||||
|
||||
/// Trusted source identity attached by an authenticated transport boundary.
|
||||
///
|
||||
/// Public clients cannot select this value directly. Runtime/Backend adapters
|
||||
/// stamp it before forwarding an accepted Submit or Notify to a Worker.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
|
||||
pub enum AuthenticatedInputSource {
|
||||
/// Assigned whenever a serialized tracked method crosses an untrusted
|
||||
/// protocol boundary. Receivers must handle it exactly like public input.
|
||||
UntrustedWire,
|
||||
Account {
|
||||
account_id: String,
|
||||
},
|
||||
Worker {
|
||||
runtime_id: String,
|
||||
worker_id: String,
|
||||
},
|
||||
SubWorker {
|
||||
session_id: String,
|
||||
},
|
||||
Backend {
|
||||
operation_id: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl Default for AuthenticatedInputSource {
|
||||
fn default() -> Self {
|
||||
Self::UntrustedWire
|
||||
}
|
||||
}
|
||||
|
||||
impl AuthenticatedInputSource {
|
||||
pub fn namespace(&self) -> String {
|
||||
match self {
|
||||
Self::UntrustedWire => "untrusted-wire".into(),
|
||||
Self::Account { account_id } => format!("account:{account_id}"),
|
||||
Self::Worker {
|
||||
runtime_id,
|
||||
worker_id,
|
||||
} => format!("worker:{runtime_id}:{worker_id}"),
|
||||
Self::SubWorker { session_id } => format!("sub_worker:{session_id}"),
|
||||
Self::Backend { operation_id } => format!("backend:{operation_id}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
#[serde(tag = "method", content = "params", rename_all = "snake_case")]
|
||||
pub enum Method {
|
||||
Run {
|
||||
/// Durably accept typed input for immediate activation or the session FIFO.
|
||||
///
|
||||
/// `submission_request_id` is generated by the authenticated caller and is
|
||||
/// used only for idempotent retry. Worker allocates the durable
|
||||
/// `submission_id` returned by [`Event::SubmissionAccepted`].
|
||||
Submit {
|
||||
submission_request_id: String,
|
||||
input: Vec<Segment>,
|
||||
},
|
||||
/// Runtime-internal Run carrying an opaque correlation id that is committed
|
||||
/// with the resulting UserInput entry. This variant is not serializable on
|
||||
/// the public Client → Worker protocol.
|
||||
#[serde(skip)]
|
||||
/// Authenticated transport form of Submit. Trusted adapters replace
|
||||
/// public Submit before forwarding it to the Worker.
|
||||
#[cfg_attr(feature = "typescript", ts(skip))]
|
||||
RunTracked {
|
||||
SubmitTracked {
|
||||
submission_request_id: String,
|
||||
input: Vec<Segment>,
|
||||
submission_id: String,
|
||||
#[serde(skip_deserializing, default)]
|
||||
source: AuthenticatedInputSource,
|
||||
},
|
||||
/// Human-readable text injected into the target Worker's LLM context
|
||||
/// as a non-blocking system message. `auto_run` controls whether an
|
||||
@@ -50,25 +114,54 @@ pub enum Method {
|
||||
/// No side effects beyond LLM context; use `WorkerEvent` for typed
|
||||
/// lifecycle reports.
|
||||
Notify {
|
||||
notification_request_id: String,
|
||||
message: String,
|
||||
#[serde(default = "default_true", skip_serializing_if = "is_true")]
|
||||
auto_run: bool,
|
||||
},
|
||||
/// Authenticated transport form of Notify.
|
||||
#[cfg_attr(feature = "typescript", ts(skip))]
|
||||
NotifyTracked {
|
||||
notification_request_id: String,
|
||||
message: String,
|
||||
#[serde(default = "default_true", skip_serializing_if = "is_true")]
|
||||
auto_run: bool,
|
||||
#[serde(skip_deserializing, default)]
|
||||
source: AuthenticatedInputSource,
|
||||
},
|
||||
/// Typed lifecycle report from a child Worker to its direct parent.
|
||||
WorkerEvent(WorkerEvent),
|
||||
/// Return the authoritative FIFO summary without exposing queued payloads.
|
||||
ListPendingSubmissions,
|
||||
/// Remove one queued submission. Running or already activated submissions
|
||||
/// are immutable and therefore cannot be cancelled here.
|
||||
CancelPendingSubmission {
|
||||
submission_id: String,
|
||||
expected_revision: u64,
|
||||
},
|
||||
/// Remove every queued submission while preserving the active run.
|
||||
ClearPendingSubmissions {
|
||||
expected_revision: u64,
|
||||
},
|
||||
/// Activate the next queued submission while the Worker is idle. This is an
|
||||
/// explicit recovery operation and never resumes a paused run implicitly.
|
||||
ContinuePending {
|
||||
expected_revision: u64,
|
||||
expected_head_id: String,
|
||||
},
|
||||
Resume,
|
||||
Cancel,
|
||||
/// Stop the in-flight turn and transition to `Paused`.
|
||||
///
|
||||
/// Unlike `Cancel` (which discards and returns to `Idle`), a paused
|
||||
/// Worker can resume the interrupted work via `Resume`, or start a
|
||||
/// fresh turn via `Run` (orphan `tool_use` items are closed with a
|
||||
/// Worker can resume the interrupted work via `Resume`, or accept a
|
||||
/// fresh `Submit` (orphan `tool_use` items are closed with a
|
||||
/// synthetic tool result before the new user message is appended).
|
||||
Pause,
|
||||
/// Request an explicit compaction while the Worker is otherwise idle.
|
||||
///
|
||||
/// This is a typed control method: clients must not send `compact` as a
|
||||
/// `Method::Run` user message.
|
||||
/// `Method::Submit` user message.
|
||||
Compact,
|
||||
/// Ask the Worker to list valid rewind targets from its authoritative session log.
|
||||
ListRewindTargets,
|
||||
@@ -181,7 +274,7 @@ impl WorkerEvent {
|
||||
|
||||
/// One typed piece of a user submission.
|
||||
///
|
||||
/// `Method::Run` and `Event::UserMessage` carry `Vec<Segment>`. Dumb
|
||||
/// `Method::Submit` and `Event::UserMessage` carry `Vec<Segment>`. Dumb
|
||||
/// clients (CLI piping, scripts) only need to produce a single
|
||||
/// `Segment::Text`; richer clients (TUI / GUI) construct typed atoms
|
||||
/// (paste chips, file refs) and
|
||||
@@ -404,12 +497,13 @@ impl Segment {
|
||||
}
|
||||
|
||||
impl Method {
|
||||
/// Convenience: a `Run` carrying a single `Segment::Text`.
|
||||
/// Convenience: a `Submit` carrying a single `Segment::Text`.
|
||||
/// Used by dumb clients, inter-Worker tools, and tests that only have
|
||||
/// a string to forward.
|
||||
pub fn run_text(s: impl Into<String>) -> Self {
|
||||
Self::Run {
|
||||
input: vec![Segment::text(s)],
|
||||
pub fn submit_text(submission_request_id: impl Into<String>, text: impl Into<String>) -> Self {
|
||||
Self::Submit {
|
||||
submission_request_id: submission_request_id.into(),
|
||||
input: vec![Segment::text(text)],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -503,6 +597,39 @@ pub enum ToolResultDisposition {
|
||||
OutcomeUnknown,
|
||||
}
|
||||
|
||||
/// Durable acceptance result for one idempotent Submit request.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SubmissionDisposition {
|
||||
Started,
|
||||
Queued,
|
||||
}
|
||||
|
||||
/// Bounded public projection of one pending submission. Payload segments and
|
||||
/// provenance remain in the session log and are intentionally not exposed.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
pub struct PendingSubmissionSummary {
|
||||
pub submission_id: String,
|
||||
pub accepted_at_ms: u64,
|
||||
pub segment_count: u32,
|
||||
pub byte_len: u64,
|
||||
}
|
||||
|
||||
/// Revisioned session-owned FIFO projection used by snapshots and live events.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
pub struct PendingSubmissionsSnapshot {
|
||||
pub revision: u64,
|
||||
#[serde(default)]
|
||||
pub notification_count: u32,
|
||||
#[serde(default)]
|
||||
pub head_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub submissions: Vec<PendingSubmissionSummary>,
|
||||
}
|
||||
|
||||
/// Canonical, storage-independent projection of committed session history.
|
||||
///
|
||||
/// Worker protocols expose this DTO instead of append-log records. New
|
||||
@@ -511,6 +638,8 @@ pub enum ToolResultDisposition {
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
pub struct SessionSnapshot {
|
||||
#[serde(default)]
|
||||
pub pending_submissions: PendingSubmissionsSnapshot,
|
||||
pub entries: Vec<SessionSnapshotEntry>,
|
||||
}
|
||||
|
||||
@@ -609,16 +738,27 @@ pub struct SessionToolAttachment {
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
#[serde(tag = "event", content = "data", rename_all = "snake_case")]
|
||||
pub enum Event {
|
||||
/// A user input message was accepted, persisted as
|
||||
/// `LogEntry::AnnotatedUserInput`, and is about to start a new turn.
|
||||
/// Broadcast to every subscribed client so TUI / GUI instances show
|
||||
/// the same user line that reconnect snapshots would replay from
|
||||
/// history; clients must not synthesize a separate pending/fake
|
||||
/// message for accepted runs.
|
||||
///
|
||||
/// Fires exactly once per committed user input, after
|
||||
/// `InvokeStart { kind: UserSend }` and before the first
|
||||
/// `TurnStart`. Rejected runs (e.g. `AlreadyRunning`) do not emit.
|
||||
/// Durable Submit acceptance. A `Started` receipt follows the atomic
|
||||
/// UserInput commit; a `Queued` receipt follows the durable FIFO checkpoint.
|
||||
/// Repeating the same request id and exact payload returns the same receipt
|
||||
/// without appending or activating twice.
|
||||
SubmissionAccepted {
|
||||
submission_request_id: String,
|
||||
submission_id: String,
|
||||
disposition: SubmissionDisposition,
|
||||
},
|
||||
/// Correlated rejection before durable acceptance.
|
||||
SubmissionRejected {
|
||||
submission_request_id: String,
|
||||
message: String,
|
||||
},
|
||||
/// Revisioned FIFO replacement following enqueue, activation, cancel, or clear.
|
||||
PendingSubmissionsChanged {
|
||||
pending: PendingSubmissionsSnapshot,
|
||||
},
|
||||
/// A user input message persisted as `LogEntry::AnnotatedUserInput` and
|
||||
/// activated for a turn. Broadcast to every subscribed client so TUI / GUI
|
||||
/// instances show the same user line that reconnect snapshots replay.
|
||||
UserMessage {
|
||||
segments: Vec<Segment>,
|
||||
},
|
||||
@@ -641,7 +781,7 @@ pub enum Event {
|
||||
///
|
||||
/// Marker event for the start of an Invoke range; the range extends
|
||||
/// implicitly until the next `InvokeStart`. Fires for every accepted
|
||||
/// `Method::Run` (kind=`UserSend`), `Method::Notify` (kind=`Notify`),
|
||||
/// `Method::Submit` (kind=`UserSend`), `Method::Notify` (kind=`Notify`),
|
||||
/// `Method::WorkerEvent` re-injection (kind=`WorkerEvent`), and any other
|
||||
/// IDLE-breaking trigger. Mid-run interrupts (e.g. hook output,
|
||||
/// typed system reminder insertion that doesn't break IDLE) do not
|
||||
@@ -1193,7 +1333,7 @@ pub enum TurnResult {
|
||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum InvokeKind {
|
||||
/// `Method::Run` — a user submission.
|
||||
/// `Method::Submit` — a user submission.
|
||||
UserSend,
|
||||
/// `Method::Notify` — free-text notification injected into history.
|
||||
Notify,
|
||||
@@ -1216,7 +1356,7 @@ pub enum RunResult {
|
||||
Finished,
|
||||
Paused,
|
||||
LimitReached,
|
||||
/// The accepted Method::Run produced no assistant/tool output before
|
||||
/// The accepted Method::Submit produced no assistant/tool output before
|
||||
/// user interruption, so the Worker rolled the submit-time turn state back
|
||||
/// to its pre-submit snapshot. Clients should treat the Worker as Idle and
|
||||
/// restore the just-submitted input into the editable composer if desired.
|
||||
@@ -1285,26 +1425,30 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn method_run_json_roundtrip() {
|
||||
let json = r#"{"method":"run","params":{"input":[{"kind":"text","content":"Hello"}]}}"#;
|
||||
fn method_submit_json_roundtrip_and_run_is_rejected() {
|
||||
let json = r#"{"method":"submit","params":{"submission_request_id":"request-1","input":[{"kind":"text","content":"Hello"}]}}"#;
|
||||
let method: Method = serde_json::from_str(json).unwrap();
|
||||
match &method {
|
||||
Method::Run { input } => {
|
||||
Method::Submit { input, .. } => {
|
||||
assert_eq!(input.len(), 1);
|
||||
match &input[0] {
|
||||
Segment::Text { content } => assert_eq!(content, "Hello"),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
other => panic!("expected Submit, got {other:?}"),
|
||||
}
|
||||
let serialized = serde_json::to_string(&method).unwrap();
|
||||
assert_eq!(serialized, json);
|
||||
assert!(
|
||||
serde_json::from_str::<Method>(r#"{"method":"run","params":{"input":[]}}"#).is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn method_run_paste_segment_roundtrip() {
|
||||
let method = Method::Run {
|
||||
fn method_submit_paste_segment_roundtrip() {
|
||||
let method = Method::Submit {
|
||||
submission_request_id: "request-1".to_string(),
|
||||
input: vec![
|
||||
Segment::text("see "),
|
||||
Segment::Paste {
|
||||
@@ -1318,7 +1462,7 @@ mod tests {
|
||||
let json = serde_json::to_string(&method).unwrap();
|
||||
let decoded: Method = serde_json::from_str(&json).unwrap();
|
||||
match decoded {
|
||||
Method::Run { input } => {
|
||||
Method::Submit { input, .. } => {
|
||||
assert_eq!(input.len(), 2);
|
||||
match &input[1] {
|
||||
Segment::Paste {
|
||||
@@ -1335,7 +1479,7 @@ mod tests {
|
||||
other => panic!("expected Paste, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
other => panic!("expected Submit, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1389,8 +1533,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn method_run_flow_segment_roundtrip() {
|
||||
let method = Method::Run {
|
||||
fn method_submit_flow_segment_roundtrip() {
|
||||
let method = Method::Submit {
|
||||
submission_request_id: "request-1".to_string(),
|
||||
input: vec![
|
||||
Segment::Flow {
|
||||
selector: "builtin:coder-review".to_string(),
|
||||
@@ -1404,7 +1549,7 @@ mod tests {
|
||||
let decoded = serde_json::from_str::<Method>(&json).unwrap();
|
||||
assert!(matches!(
|
||||
decoded,
|
||||
Method::Run { input }
|
||||
Method::Submit { input, .. }
|
||||
if matches!(
|
||||
input.as_slice(),
|
||||
[
|
||||
@@ -1416,15 +1561,26 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_tracked_run_is_not_public_protocol_json() {
|
||||
let method = Method::RunTracked {
|
||||
fn authenticated_submit_replaces_wire_source_with_transport_identity() {
|
||||
let method = Method::SubmitTracked {
|
||||
input: vec![Segment::text("private")],
|
||||
submission_id: "submission-1".to_string(),
|
||||
submission_request_id: "request-1".to_string(),
|
||||
source: AuthenticatedInputSource::Account {
|
||||
account_id: "account-1".into(),
|
||||
},
|
||||
};
|
||||
assert!(serde_json::to_string(&method).is_err());
|
||||
let json = serde_json::to_string(&method).unwrap();
|
||||
let decoded = serde_json::from_str::<Method>(&json).unwrap();
|
||||
assert!(matches!(
|
||||
decoded,
|
||||
Method::SubmitTracked {
|
||||
source: AuthenticatedInputSource::UntrustedWire,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert!(
|
||||
serde_json::from_str::<Method>(
|
||||
r#"{"method":"run_tracked","input":[],"submission_id":"forged"}"#,
|
||||
r#"{"method":"submit_tracked","input":[],"submission_request_id":"forged"}"#,
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
@@ -1442,16 +1598,16 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn method_run_with_unknown_segment_decodes() {
|
||||
let json = r#"{"method":"run","params":{"input":[{"kind":"text","content":"hi"},{"kind":"future_thing","x":1}]}}"#;
|
||||
fn method_submit_with_unknown_segment_decodes() {
|
||||
let json = r#"{"method":"submit","params":{"submission_request_id":"request-1","input":[{"kind":"text","content":"hi"},{"kind":"future_thing","x":1}]}}"#;
|
||||
let method: Method = serde_json::from_str(json).unwrap();
|
||||
match method {
|
||||
Method::Run { input } => {
|
||||
Method::Submit { input, .. } => {
|
||||
assert_eq!(input.len(), 2);
|
||||
assert!(matches!(input[0], Segment::Text { .. }));
|
||||
assert!(matches!(input[1], Segment::Unknown));
|
||||
}
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
other => panic!("expected Submit, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1648,11 +1804,11 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn method_notify_json_roundtrip_defaults_to_auto_run() {
|
||||
let json = r#"{"method":"notify","params":{"message":"turn done"}}"#;
|
||||
let json = r#"{"method":"notify","params":{"notification_request_id":"notification-1","message":"turn done"}}"#;
|
||||
let method: Method = serde_json::from_str(json).unwrap();
|
||||
assert!(matches!(
|
||||
method,
|
||||
Method::Notify { ref message, auto_run: true } if message == "turn done"
|
||||
Method::Notify { ref message, auto_run: true, .. } if message == "turn done"
|
||||
));
|
||||
let serialized = serde_json::to_string(&method).unwrap();
|
||||
assert_eq!(serialized, json);
|
||||
@@ -1660,11 +1816,11 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn method_notify_weak_json_roundtrip_serializes_auto_run_false() {
|
||||
let json = r#"{"method":"notify","params":{"message":"progress","auto_run":false}}"#;
|
||||
let json = r#"{"method":"notify","params":{"notification_request_id":"notification-1","message":"progress","auto_run":false}}"#;
|
||||
let method: Method = serde_json::from_str(json).unwrap();
|
||||
assert!(matches!(
|
||||
method,
|
||||
Method::Notify { ref message, auto_run: false } if message == "progress"
|
||||
Method::Notify { ref message, auto_run: false, .. } if message == "progress"
|
||||
));
|
||||
assert_eq!(serde_json::to_string(&method).unwrap(), json);
|
||||
}
|
||||
@@ -1725,6 +1881,7 @@ mod tests {
|
||||
fn event_snapshot_format() {
|
||||
let event = Event::Snapshot {
|
||||
session: SessionSnapshot {
|
||||
pending_submissions: PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![SessionSnapshotEntry {
|
||||
entry_id: "entry-1".into(),
|
||||
timestamp: 1,
|
||||
@@ -1776,6 +1933,7 @@ mod tests {
|
||||
|
||||
let event = Event::Snapshot {
|
||||
session: SessionSnapshot {
|
||||
pending_submissions: PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: Greeting {
|
||||
@@ -1844,6 +2002,7 @@ mod tests {
|
||||
fn event_segment_rotated_roundtrip() {
|
||||
let event = Event::SegmentRotated {
|
||||
session: SessionSnapshot {
|
||||
pending_submissions: PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
};
|
||||
|
||||
@@ -8,11 +8,11 @@ use crate::{
|
||||
CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot,
|
||||
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
|
||||
InvokeKind, MemoryWorkerEvent, Method, PasteArtifactAvailability, PasteArtifactMediaType,
|
||||
PasteArtifactRef, Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult,
|
||||
ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole,
|
||||
SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment,
|
||||
ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent,
|
||||
WorkerStatus,
|
||||
PasteArtifactRef, PendingSubmissionSummary, PendingSubmissionsSnapshot, Permission,
|
||||
RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart,
|
||||
SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry,
|
||||
SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, ToolResultDisposition,
|
||||
TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent, WorkerStatus,
|
||||
subscription::{
|
||||
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
||||
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
|
||||
@@ -75,6 +75,9 @@ pub fn generated_protocol_types() -> String {
|
||||
push_decl::<SessionToolAttachment>(&cfg, &mut output);
|
||||
push_decl::<SessionSnapshotEntryData>(&cfg, &mut output);
|
||||
push_decl::<SessionSnapshotEntry>(&cfg, &mut output);
|
||||
push_decl::<PendingSubmissionSummary>(&cfg, &mut output);
|
||||
push_decl::<PendingSubmissionsSnapshot>(&cfg, &mut output);
|
||||
push_decl::<SubmissionDisposition>(&cfg, &mut output);
|
||||
push_decl::<SessionSnapshot>(&cfg, &mut output);
|
||||
push_decl::<InternalWorkerKind>(&cfg, &mut output);
|
||||
push_decl::<InternalWorkerRef>(&cfg, &mut output);
|
||||
|
||||
@@ -21,8 +21,10 @@ use crate::segment_log::LogEntry;
|
||||
use crate::store::{Store, StoreError};
|
||||
use crate::uploaded_file::{
|
||||
bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files,
|
||||
delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs,
|
||||
read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file,
|
||||
delete_uncommitted_uploaded_files, delete_uploaded_file, finalize_uploaded_file_binding,
|
||||
list_uploaded_file_refs, pin_uploaded_file, read_uploaded_file, read_uploaded_file_by_id,
|
||||
reconcile_uploaded_file_pins, release_uploaded_file_pin, uploaded_file_has_pending_owner,
|
||||
write_uploaded_file,
|
||||
};
|
||||
use crate::{
|
||||
PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext,
|
||||
@@ -518,6 +520,61 @@ impl Store for FsStore {
|
||||
}
|
||||
}
|
||||
|
||||
fn pin_uploaded_file(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
reference: &UploadedFileRef,
|
||||
owner_id: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
let _guard = self
|
||||
.append_lock
|
||||
.lock()
|
||||
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
||||
pin_uploaded_file(&self.paste_artifact_dir(session_id), reference, owner_id)
|
||||
}
|
||||
|
||||
fn release_uploaded_file_pin(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
artifact_id: &str,
|
||||
owner_id: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
let _guard = self
|
||||
.append_lock
|
||||
.lock()
|
||||
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
||||
release_uploaded_file_pin(&self.paste_artifact_dir(session_id), artifact_id, owner_id)
|
||||
}
|
||||
|
||||
fn finalize_uploaded_file_binding(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
artifact_id: &str,
|
||||
source_entry_id: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
let _guard = self
|
||||
.append_lock
|
||||
.lock()
|
||||
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
||||
finalize_uploaded_file_binding(
|
||||
&self.paste_artifact_dir(session_id),
|
||||
artifact_id,
|
||||
source_entry_id,
|
||||
)
|
||||
}
|
||||
|
||||
fn reconcile_uploaded_file_pins(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
live_owner_ids: &[String],
|
||||
) -> Result<u64, StoreError> {
|
||||
let _guard = self
|
||||
.append_lock
|
||||
.lock()
|
||||
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
||||
reconcile_uploaded_file_pins(&self.paste_artifact_dir(session_id), live_owner_ids)
|
||||
}
|
||||
|
||||
fn delete_uploaded_file(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
@@ -541,13 +598,18 @@ impl Store for FsStore {
|
||||
let Some(source_entry_id) = reference.source_entry_id.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
if !self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? {
|
||||
clear_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
|
||||
if delete_uploaded_file(&dir, &reference.artifact_id)? {
|
||||
removed = removed
|
||||
.checked_add(1)
|
||||
.ok_or(StoreError::ArtifactQuotaExceeded)?;
|
||||
}
|
||||
if self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? {
|
||||
finalize_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
|
||||
continue;
|
||||
}
|
||||
if uploaded_file_has_pending_owner(&dir, &reference.artifact_id)? {
|
||||
continue;
|
||||
}
|
||||
clear_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
|
||||
if delete_uploaded_file(&dir, &reference.artifact_id)? {
|
||||
removed = removed
|
||||
.checked_add(1)
|
||||
.ok_or(StoreError::ArtifactQuotaExceeded)?;
|
||||
}
|
||||
}
|
||||
Ok(removed)
|
||||
@@ -865,6 +927,106 @@ mod tests {
|
||||
assert!(store.read_uploaded_file(owner, &reference).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_upload_pin_survives_cleanup_until_release_or_history_binding() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let store = FsStore::new(tmp.path()).unwrap();
|
||||
let session_id = new_session_id();
|
||||
let limits = UploadedFileLimits {
|
||||
max_file_bytes: 64,
|
||||
max_session_bytes: 128,
|
||||
};
|
||||
let pending = store
|
||||
.write_uploaded_file(session_id, "pending.txt", "text/plain", b"pending", limits)
|
||||
.unwrap();
|
||||
store
|
||||
.pin_uploaded_file(session_id, &pending, "submission-1")
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
store.pin_uploaded_file(session_id, &pending, "submission-other"),
|
||||
Err(StoreError::ArtifactAlreadyCommitted)
|
||||
));
|
||||
drop(store);
|
||||
let store = FsStore::new(tmp.path()).unwrap();
|
||||
assert_eq!(
|
||||
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.read_uploaded_file_by_id(session_id, &pending.artifact_id)
|
||||
.unwrap()
|
||||
.1,
|
||||
b"pending"
|
||||
);
|
||||
|
||||
let fork_session_id = new_session_id();
|
||||
assert_eq!(
|
||||
store
|
||||
.copy_committed_uploaded_files(session_id, fork_session_id)
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.read_uploaded_file_by_id(fork_session_id, &pending.artifact_id)
|
||||
.is_err()
|
||||
);
|
||||
|
||||
let committed = store
|
||||
.bind_uploaded_file(session_id, &pending, "entry-1")
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
|
||||
0
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.read_uploaded_file_by_id(session_id, &pending.artifact_id)
|
||||
.is_ok()
|
||||
);
|
||||
store
|
||||
.create_segment(
|
||||
session_id,
|
||||
new_segment_id(),
|
||||
&[LogEntry::InputSegmentsCheckpoint {
|
||||
ts: 1,
|
||||
user_segments: vec![vec![protocol::Segment::UploadedFile {
|
||||
file: committed.clone(),
|
||||
}]],
|
||||
}],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
|
||||
0
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.release_uploaded_file_pin(session_id, &pending.artifact_id, "submission-1")
|
||||
.is_err()
|
||||
);
|
||||
|
||||
let releasable = store
|
||||
.write_uploaded_file(session_id, "cancelled.txt", "text/plain", b"cancel", limits)
|
||||
.unwrap();
|
||||
store
|
||||
.pin_uploaded_file(session_id, &releasable, "submission-2")
|
||||
.unwrap();
|
||||
store
|
||||
.release_uploaded_file_pin(session_id, &releasable.artifact_id, "submission-2")
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
|
||||
1
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.read_uploaded_file_by_id(session_id, &releasable.artifact_id)
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uploaded_file_validation_and_shared_quota_fail_closed() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
|
||||
@@ -183,6 +183,7 @@ fn canonicalize_history_entry(
|
||||
item,
|
||||
metadata: legacy_metadata(segment_id, line_index, 0),
|
||||
},
|
||||
extensions: Vec::new(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -71,7 +71,7 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
|
||||
entries.push(history_entry(entry, *ts, data));
|
||||
}
|
||||
}
|
||||
LogEntry::AnnotatedSystemItem { ts, entry } => entries.push(system_entry(
|
||||
LogEntry::AnnotatedSystemItem { ts, entry, .. } => entries.push(system_entry(
|
||||
&entry.item,
|
||||
entry.metadata.entry_id.0.clone(),
|
||||
*ts,
|
||||
@@ -100,7 +100,10 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
|
||||
}
|
||||
}
|
||||
|
||||
SessionSnapshot { entries }
|
||||
SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries,
|
||||
}
|
||||
}
|
||||
|
||||
fn extend_history(
|
||||
|
||||
@@ -287,6 +287,7 @@ pub fn append_system_item(
|
||||
LogEntry::AnnotatedSystemItem {
|
||||
ts: segment_log::now_millis(),
|
||||
entry,
|
||||
extensions: Vec::new(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@@ -112,6 +112,8 @@ pub enum LogEntry {
|
||||
AnnotatedSystemItem {
|
||||
ts: u64,
|
||||
entry: LoggedSystemHistoryEntry,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
extensions: Vec<SessionExtension>,
|
||||
},
|
||||
|
||||
/// Turn boundary. Records the turn count after increment.
|
||||
@@ -312,12 +314,19 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
|
||||
state.annotated_history.push(entry.clone());
|
||||
state.history.push(Item::from(entry.item.clone()));
|
||||
}
|
||||
LogEntry::AnnotatedSystemItem { entry, .. } => {
|
||||
LogEntry::AnnotatedSystemItem {
|
||||
entry, extensions, ..
|
||||
} => {
|
||||
state.annotated_history.push(LoggedHistoryEntry {
|
||||
item: LoggedItem::from(entry.item.to_history_item()),
|
||||
metadata: entry.metadata.clone(),
|
||||
});
|
||||
state.history.push(entry.item.to_history_item());
|
||||
state.extensions.extend(
|
||||
extensions
|
||||
.iter()
|
||||
.map(|extension| (extension.domain.clone(), extension.payload.clone())),
|
||||
);
|
||||
}
|
||||
LogEntry::TurnEnd { turn_count, .. } => {
|
||||
if let Some(active_turn_count) = &mut state.active_run_turn_count {
|
||||
|
||||
@@ -226,6 +226,47 @@ pub trait Store: Send + Sync {
|
||||
Err(StoreError::PasteArtifactUnsupported)
|
||||
}
|
||||
|
||||
/// Retain an uploaded file while a durable pending operation owns it.
|
||||
fn pin_uploaded_file(
|
||||
&self,
|
||||
_session_id: SessionId,
|
||||
_reference: &UploadedFileRef,
|
||||
_owner_id: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
Err(StoreError::PasteArtifactUnsupported)
|
||||
}
|
||||
|
||||
/// Release a pending-operation pin without changing committed ownership.
|
||||
fn release_uploaded_file_pin(
|
||||
&self,
|
||||
_session_id: SessionId,
|
||||
_artifact_id: &str,
|
||||
_owner_id: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
Err(StoreError::PasteArtifactUnsupported)
|
||||
}
|
||||
|
||||
/// Complete the pending-to-history handoff after the history entry commits.
|
||||
fn finalize_uploaded_file_binding(
|
||||
&self,
|
||||
_session_id: SessionId,
|
||||
_artifact_id: &str,
|
||||
_source_entry_id: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
Err(StoreError::PasteArtifactUnsupported)
|
||||
}
|
||||
|
||||
/// Clear pending-operation pins that have no owner in restored durable
|
||||
/// Worker Session state. This repairs an interrupted pin-before-checkpoint
|
||||
/// acceptance without disturbing live queue owners or committed history.
|
||||
fn reconcile_uploaded_file_pins(
|
||||
&self,
|
||||
_session_id: SessionId,
|
||||
_live_owner_ids: &[String],
|
||||
) -> Result<u64, StoreError> {
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
/// Delete an uncommitted uploaded file owned by `session_id`.
|
||||
fn delete_uploaded_file(
|
||||
&self,
|
||||
|
||||
@@ -24,6 +24,12 @@ pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8;
|
||||
pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256;
|
||||
const MAX_FILE_NAME_CHARS: usize = 255;
|
||||
const MAX_MEDIA_TYPE_BYTES: usize = 127;
|
||||
fn validate_pending_owner_id(owner_id: &str) -> Result<()> {
|
||||
if owner_id.is_empty() || owner_id.len() > 256 {
|
||||
return Err(StoreError::ArtifactIntegrityMismatch);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct UploadedFileLimits {
|
||||
@@ -59,6 +65,8 @@ struct StoredUploadedFile {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
source_entry_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pending_owner_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
upload_context: Option<UploadedFileUploadContext>,
|
||||
content_base64: String,
|
||||
}
|
||||
@@ -291,6 +299,7 @@ pub(crate) fn write_uploaded_file(
|
||||
byte_len,
|
||||
sha256: sha256.clone(),
|
||||
source_entry_id: None,
|
||||
pending_owner_id: None,
|
||||
upload_context: context.cloned(),
|
||||
content_base64: BASE64.encode(content),
|
||||
};
|
||||
@@ -338,6 +347,12 @@ pub(crate) fn read_uploaded_file_by_id(
|
||||
Ok((reference, content))
|
||||
}
|
||||
|
||||
pub(crate) fn uploaded_file_has_pending_owner(dir: &Path, artifact_id: &str) -> Result<bool> {
|
||||
let path = record_path(dir, artifact_id)?;
|
||||
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?;
|
||||
Ok(stored.pending_owner_id.is_some())
|
||||
}
|
||||
|
||||
pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Result<Vec<u8>> {
|
||||
let (stored_reference, content) = read_uploaded_file_by_id(dir, &reference.artifact_id)?;
|
||||
if stored_reference.file_name != reference.file_name
|
||||
@@ -376,6 +391,98 @@ pub(crate) fn clear_uploaded_file_binding(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn pin_uploaded_file(
|
||||
dir: &Path,
|
||||
reference: &UploadedFileRef,
|
||||
owner_id: &str,
|
||||
) -> Result<()> {
|
||||
validate_pending_owner_id(owner_id)?;
|
||||
if reference.source_entry_id.is_some() {
|
||||
return Err(StoreError::ArtifactAlreadyCommitted);
|
||||
}
|
||||
let aggregate_lock = fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.open(dir.join(".aggregate.lock"))?;
|
||||
FileExt::lock_exclusive(&aggregate_lock)?;
|
||||
let path = record_path(dir, &reference.artifact_id)?;
|
||||
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||
if stored.file_name != reference.file_name
|
||||
|| stored.media_type != reference.media_type
|
||||
|| stored.created_at_ms != reference.created_at_ms
|
||||
|| stored.byte_len != reference.byte_len
|
||||
|| stored.sha256 != reference.sha256
|
||||
{
|
||||
return Err(StoreError::ArtifactIntegrityMismatch);
|
||||
}
|
||||
if stored.source_entry_id.is_some() {
|
||||
return Err(StoreError::ArtifactAlreadyCommitted);
|
||||
}
|
||||
if let Some(existing_owner) = stored.pending_owner_id.as_deref() {
|
||||
return if existing_owner == owner_id {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(StoreError::ArtifactAlreadyCommitted)
|
||||
};
|
||||
}
|
||||
stored.pending_owner_id = Some(owner_id.to_owned());
|
||||
let temp = dir.join(format!(".{}.file.pin.tmp", reference.artifact_id));
|
||||
fs::write(&temp, serde_json::to_vec(&stored)?)?;
|
||||
fs::rename(temp, path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn release_uploaded_file_pin(
|
||||
dir: &Path,
|
||||
artifact_id: &str,
|
||||
owner_id: &str,
|
||||
) -> Result<()> {
|
||||
validate_pending_owner_id(owner_id)?;
|
||||
let aggregate_lock = fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.open(dir.join(".aggregate.lock"))?;
|
||||
FileExt::lock_exclusive(&aggregate_lock)?;
|
||||
let path = record_path(dir, artifact_id)?;
|
||||
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||
if stored.pending_owner_id.as_deref() != Some(owner_id) {
|
||||
return Err(StoreError::ArtifactIntegrityMismatch);
|
||||
}
|
||||
stored.pending_owner_id = None;
|
||||
let temp = dir.join(format!(".{artifact_id}.file.unpin.tmp"));
|
||||
fs::write(&temp, serde_json::to_vec(&stored)?)?;
|
||||
fs::rename(temp, path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn finalize_uploaded_file_binding(
|
||||
dir: &Path,
|
||||
artifact_id: &str,
|
||||
source_entry_id: &str,
|
||||
) -> Result<()> {
|
||||
let aggregate_lock = fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.open(dir.join(".aggregate.lock"))?;
|
||||
FileExt::lock_exclusive(&aggregate_lock)?;
|
||||
let path = record_path(dir, artifact_id)?;
|
||||
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||
if stored.source_entry_id.as_deref() != Some(source_entry_id) {
|
||||
return Err(StoreError::ArtifactIntegrityMismatch);
|
||||
}
|
||||
if stored.pending_owner_id.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
stored.pending_owner_id = None;
|
||||
let temp = dir.join(format!(".{artifact_id}.file.finalize.tmp"));
|
||||
fs::write(&temp, serde_json::to_vec(&stored)?)?;
|
||||
fs::rename(temp, path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn bind_uploaded_file(
|
||||
dir: &Path,
|
||||
reference: &UploadedFileRef,
|
||||
@@ -479,6 +586,40 @@ pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path
|
||||
Ok(copied)
|
||||
}
|
||||
|
||||
pub(crate) fn reconcile_uploaded_file_pins(dir: &Path, live_owner_ids: &[String]) -> Result<u64> {
|
||||
fs::create_dir_all(dir)?;
|
||||
let aggregate_lock = fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.open(dir.join(".aggregate.lock"))?;
|
||||
FileExt::lock_exclusive(&aggregate_lock)?;
|
||||
let mut reconciled = 0_u64;
|
||||
for entry in fs::read_dir(dir)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
let Some(artifact_id) = file_name.strip_suffix(".file.json") else {
|
||||
continue;
|
||||
};
|
||||
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||
let Some(owner_id) = stored.pending_owner_id.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
if live_owner_ids.iter().any(|live| live == owner_id) {
|
||||
continue;
|
||||
}
|
||||
stored.pending_owner_id = None;
|
||||
let temp = dir.join(format!(".{artifact_id}.file.reconcile.tmp"));
|
||||
fs::write(&temp, serde_json::to_vec(&stored)?)?;
|
||||
fs::rename(temp, path)?;
|
||||
reconciled = reconciled.saturating_add(1);
|
||||
}
|
||||
Ok(reconciled)
|
||||
}
|
||||
|
||||
pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
|
||||
fs::create_dir_all(dir)?;
|
||||
let aggregate_lock = fs::OpenOptions::new()
|
||||
@@ -499,7 +640,7 @@ pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
|
||||
continue;
|
||||
}
|
||||
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||
if stored.source_entry_id.is_none() {
|
||||
if stored.source_entry_id.is_none() && stored.pending_owner_id.is_none() {
|
||||
fs::remove_file(path)?;
|
||||
removed = removed
|
||||
.checked_add(1)
|
||||
@@ -523,7 +664,7 @@ pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result<bool
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
||||
Err(error) => return Err(error.into()),
|
||||
};
|
||||
if stored.source_entry_id.is_some() {
|
||||
if stored.source_entry_id.is_some() || stored.pending_owner_id.is_some() {
|
||||
return Err(StoreError::ArtifactAlreadyCommitted);
|
||||
}
|
||||
match fs::remove_file(path) {
|
||||
|
||||
@@ -99,7 +99,10 @@ async fn in_process_host_runs_text_and_read_tool_then_shuts_down() {
|
||||
let mut protocol_client = host.connect();
|
||||
|
||||
protocol_client
|
||||
.send(&Method::run_text("read the probe"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"read the probe",
|
||||
))
|
||||
.await
|
||||
.expect("submit input");
|
||||
|
||||
@@ -336,11 +339,15 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
|
||||
let worker_id = host.worker_id();
|
||||
let mut protocol_client = host.connect();
|
||||
protocol_client
|
||||
.send(&Method::run_text("first request"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"first request",
|
||||
))
|
||||
.await?;
|
||||
wait_for_run_end(&mut protocol_client).await?;
|
||||
protocol_client
|
||||
.send(&Method::Notify {
|
||||
notification_request_id: protocol::new_submission_request_id(),
|
||||
message: "persisted notification".to_string(),
|
||||
auto_run: true,
|
||||
})
|
||||
@@ -394,7 +401,10 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
|
||||
assert!(snapshot.contains("persisted notification"), "{snapshot}");
|
||||
|
||||
protocol_client
|
||||
.send(&Method::run_text("continue after restore"))
|
||||
.send(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"continue after restore",
|
||||
))
|
||||
.await?;
|
||||
wait_for_run_end(&mut protocol_client).await?;
|
||||
let request = second_inspection
|
||||
|
||||
+125
-154
@@ -102,23 +102,6 @@ struct RollbackSubmitState {
|
||||
turn_before: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QueuedInput {
|
||||
segments: Vec<Segment>,
|
||||
preview: String,
|
||||
}
|
||||
|
||||
impl QueuedInput {
|
||||
fn new(segments: Vec<Segment>) -> Self {
|
||||
let preview = Segment::flatten_to_text(&segments);
|
||||
Self { segments, preview }
|
||||
}
|
||||
|
||||
pub fn preview(&self) -> &str {
|
||||
&self.preview
|
||||
}
|
||||
}
|
||||
|
||||
struct ComposerInputHistory {
|
||||
entries: VecDeque<Vec<Segment>>,
|
||||
browse: Option<ComposerInputHistoryBrowse>,
|
||||
@@ -272,7 +255,7 @@ pub struct App {
|
||||
/// Current transient actionbar notice. Notices are local UI state only:
|
||||
/// they are never appended to transcript/session history or LLM context.
|
||||
actionbar_notice: Option<ActionbarNotice>,
|
||||
/// Normal composer input that is submitted as `Method::Run`.
|
||||
/// Normal composer input that is submitted as `Method::Submit`.
|
||||
pub input: InputBuffer,
|
||||
/// Separate command-line input. It is never submitted as a user message.
|
||||
pub command_input: InputBuffer,
|
||||
@@ -333,9 +316,8 @@ pub struct App {
|
||||
/// Top entry index of the task pane's visible window. Clamped on
|
||||
/// render so it never points past the end of the list.
|
||||
pub task_pane_scroll: usize,
|
||||
/// TUI-local FIFO of user inputs submitted while the Worker is already running.
|
||||
/// Entries have not been sent to the Worker yet, so they remain editable/cancellable locally.
|
||||
queued_inputs: VecDeque<QueuedInput>,
|
||||
/// Authoritative WorkerSession FIFO summary received from snapshot/live events.
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot,
|
||||
/// TUI-local readline-style composer input history. This is intentionally
|
||||
/// client-side only: recalled entries are plain drafts until submitted again.
|
||||
input_history: ComposerInputHistory,
|
||||
@@ -395,7 +377,7 @@ impl App {
|
||||
text_selection: TextSelectionState::default(),
|
||||
task_pane_open: false,
|
||||
task_pane_scroll: 0,
|
||||
queued_inputs: VecDeque::new(),
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
input_history: ComposerInputHistory::new(),
|
||||
input_history_store: None,
|
||||
pending_submit_rollback: None,
|
||||
@@ -768,18 +750,34 @@ impl App {
|
||||
return None;
|
||||
}
|
||||
self.record_input_history(segments.clone());
|
||||
if self.running {
|
||||
self.queued_inputs.push_back(QueuedInput::new(segments));
|
||||
self.input.clear();
|
||||
self.completion = None;
|
||||
return None;
|
||||
}
|
||||
self.input.clear();
|
||||
Some(self.method_for_run(segments))
|
||||
}
|
||||
|
||||
pub fn submit_notify_input(&mut self) -> Option<Method> {
|
||||
let segments = self.input.submit_segments();
|
||||
if segments_are_blank(&segments) {
|
||||
return None;
|
||||
}
|
||||
if segments
|
||||
.iter()
|
||||
.any(|segment| matches!(segment, Segment::UploadedFile { .. }))
|
||||
{
|
||||
self.push_error("Notify accepts text only; remove attachments or queue a Submit.");
|
||||
return None;
|
||||
}
|
||||
let message = Segment::flatten_to_text(&segments);
|
||||
self.record_input_history(segments);
|
||||
self.input.clear();
|
||||
Some(Method::Notify {
|
||||
notification_request_id: protocol::new_submission_request_id(),
|
||||
message,
|
||||
auto_run: true,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn restore_unsent_run(&mut self, method: &Method) {
|
||||
let Method::Run { input } = method else {
|
||||
let Method::Submit { input, .. } = method else {
|
||||
return;
|
||||
};
|
||||
self.pending_submit_rollback = None;
|
||||
@@ -787,8 +785,9 @@ impl App {
|
||||
self.input.replace_with_segments(input);
|
||||
self.completion = None;
|
||||
} else {
|
||||
self.queued_inputs
|
||||
.push_front(QueuedInput::new(input.clone()));
|
||||
self.push_error(
|
||||
"Submit transport failed; current Composer was preserved and the unsent input was not queued.",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -804,7 +803,10 @@ impl App {
|
||||
block_start: self.blocks.len(),
|
||||
turn_before: self.turn_index,
|
||||
});
|
||||
Method::Run { input: segments }
|
||||
Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: segments,
|
||||
}
|
||||
}
|
||||
|
||||
fn record_input_history(&mut self, segments: Vec<Segment>) {
|
||||
@@ -825,7 +827,7 @@ impl App {
|
||||
}
|
||||
|
||||
pub fn queued_input_count(&self) -> usize {
|
||||
self.queued_inputs.len()
|
||||
self.pending_submissions.submissions.len()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -910,36 +912,31 @@ impl App {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn continue_pending_method(&self) -> Option<Method> {
|
||||
Some(Method::ContinuePending {
|
||||
expected_revision: self.pending_submissions.revision,
|
||||
expected_head_id: self.pending_submissions.head_id.clone()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn clear_pending_method(&self) -> Method {
|
||||
Method::ClearPendingSubmissions {
|
||||
expected_revision: self.pending_submissions.revision,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancel_pending_method(&self, submission_id: String) -> Method {
|
||||
Method::CancelPendingSubmission {
|
||||
submission_id,
|
||||
expected_revision: self.pending_submissions.revision,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn next_queued_input_preview(&self) -> Option<&str> {
|
||||
self.queued_inputs.front().map(QueuedInput::preview)
|
||||
}
|
||||
|
||||
pub fn clear_queued_inputs(&mut self) -> usize {
|
||||
let cleared = self.queued_inputs.len();
|
||||
self.queued_inputs.clear();
|
||||
cleared
|
||||
}
|
||||
|
||||
pub fn restore_next_queued_input_to_composer(&mut self) -> bool {
|
||||
if self.queued_inputs.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if !self.input.is_empty() {
|
||||
self.push_error("Composer is not empty; clear it before editing queued input.");
|
||||
return false;
|
||||
}
|
||||
let Some(queued) = self.queued_inputs.pop_front() else {
|
||||
return false;
|
||||
};
|
||||
self.input_history.cancel_browse();
|
||||
self.input.replace_with_segments(&queued.segments);
|
||||
self.completion = None;
|
||||
true
|
||||
}
|
||||
|
||||
fn pop_next_queued_run(&mut self) -> Option<Method> {
|
||||
let queued = self.queued_inputs.pop_front()?;
|
||||
Some(self.method_for_run(queued.segments))
|
||||
self.pending_submissions
|
||||
.submissions
|
||||
.first()
|
||||
.map(|submission| submission.submission_id.as_str())
|
||||
}
|
||||
|
||||
pub fn clear_actionbar_notice(&mut self) {
|
||||
@@ -1123,6 +1120,11 @@ impl App {
|
||||
}
|
||||
|
||||
match event {
|
||||
Event::SubmissionAccepted { .. } => {}
|
||||
Event::SubmissionRejected { message, .. } => self.push_error(message),
|
||||
Event::PendingSubmissionsChanged { pending } => {
|
||||
self.pending_submissions = pending;
|
||||
}
|
||||
Event::UserMessage { segments } => {
|
||||
self.turn_index += 1;
|
||||
self.blocks.push(Block::TurnHeader {
|
||||
@@ -1372,9 +1374,6 @@ impl App {
|
||||
WorkerStatus::Idle
|
||||
}
|
||||
});
|
||||
if matches!(result, RunResult::Finished | RunResult::LimitReached) {
|
||||
return self.pop_next_queued_run();
|
||||
}
|
||||
}
|
||||
}
|
||||
Event::CompactStart { .. } => {
|
||||
@@ -1449,6 +1448,7 @@ impl App {
|
||||
internal_workers,
|
||||
} => {
|
||||
self.rewind_refresh_fence = false;
|
||||
self.pending_submissions = session.pending_submissions.clone();
|
||||
self.restore_snapshot(&session, greeting, in_flight);
|
||||
self.replace_internal_worker_snapshots(internal_workers);
|
||||
self.set_worker_status(status);
|
||||
@@ -2681,7 +2681,10 @@ mod rewind_refresh_tests {
|
||||
});
|
||||
|
||||
app.handle_worker_event(Event::RewindApplied {
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
input: vec![Segment::text("selected rewind input")],
|
||||
summary: summary(3),
|
||||
});
|
||||
@@ -2700,7 +2703,10 @@ mod rewind_refresh_tests {
|
||||
});
|
||||
|
||||
app.handle_worker_event(Event::RewindApplied {
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
input: vec![Segment::text("rewound input")],
|
||||
summary: summary(1),
|
||||
});
|
||||
@@ -2743,7 +2749,10 @@ mod rewind_refresh_tests {
|
||||
});
|
||||
|
||||
app.handle_worker_event(Event::RewindApplied {
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
input: vec![Segment::text("rewound input")],
|
||||
summary: summary(2),
|
||||
});
|
||||
@@ -2877,7 +2886,7 @@ mod composer_history_persistence_tests {
|
||||
path: "src/lib.rs".into(),
|
||||
},
|
||||
]);
|
||||
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
|
||||
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
|
||||
|
||||
let mut reloaded = App::new_with_input_history_store("test".into(), store);
|
||||
assert!(reloaded.browse_input_history_older());
|
||||
@@ -2958,7 +2967,7 @@ mod composer_history_persistence_tests {
|
||||
app.insert_char(c);
|
||||
}
|
||||
match app.submit_input() {
|
||||
Some(Method::Run { input }) => input,
|
||||
Some(Method::Submit { input, .. }) => input,
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
}
|
||||
}
|
||||
@@ -3424,72 +3433,44 @@ mod completion_flow_tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn running_submit_is_queued_locally_and_clears_composer() {
|
||||
fn running_submit_is_sent_to_the_worker_and_not_queued_locally() {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
insert_text(&mut app, "queued turn");
|
||||
|
||||
assert!(app.submit_input().is_none());
|
||||
let method = app.submit_input();
|
||||
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
assert_eq!(app.next_queued_input_preview(), Some("queued turn"));
|
||||
assert!(matches!(method, Some(Method::Submit { .. })));
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
assert_eq!(input_text(&app), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finished_run_auto_sends_next_queued_input() {
|
||||
fn pending_submission_projection_is_worker_authoritative() {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
insert_text(&mut app, "next turn");
|
||||
assert!(app.submit_input().is_none());
|
||||
|
||||
let method = app.handle_worker_event(Event::RunEnd {
|
||||
result: RunResult::Finished,
|
||||
app.handle_worker_event(Event::PendingSubmissionsChanged {
|
||||
pending: protocol::PendingSubmissionsSnapshot {
|
||||
revision: 3,
|
||||
notification_count: 0,
|
||||
head_id: Some("submission-1".into()),
|
||||
submissions: vec![protocol::PendingSubmissionSummary {
|
||||
submission_id: "submission-1".into(),
|
||||
accepted_at_ms: 7,
|
||||
segment_count: 2,
|
||||
byte_len: 42,
|
||||
}],
|
||||
},
|
||||
});
|
||||
|
||||
match method {
|
||||
Some(Method::Run { input }) => {
|
||||
assert_eq!(Segment::flatten_to_text(&input), "next turn");
|
||||
}
|
||||
other => panic!("expected queued Run, got {other:?}"),
|
||||
}
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limit_reached_run_auto_sends_next_queued_input() {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
insert_text(&mut app, "next after limit");
|
||||
assert!(app.submit_input().is_none());
|
||||
|
||||
let method = app.handle_worker_event(Event::RunEnd {
|
||||
result: RunResult::LimitReached,
|
||||
});
|
||||
|
||||
match method {
|
||||
Some(Method::Run { input }) => {
|
||||
assert_eq!(Segment::flatten_to_text(&input), "next after limit");
|
||||
}
|
||||
other => panic!("expected queued Run, got {other:?}"),
|
||||
}
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn paused_and_rolled_back_run_do_not_auto_send_queue() {
|
||||
for result in [RunResult::Paused, RunResult::RolledBack] {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
insert_text(&mut app, "held turn");
|
||||
assert!(app.submit_input().is_none());
|
||||
|
||||
let method = app.handle_worker_event(Event::RunEnd { result });
|
||||
|
||||
assert!(method.is_none());
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
assert_eq!(app.next_queued_input_preview(), Some("held turn"));
|
||||
}
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
assert_eq!(app.next_queued_input_preview(), Some("submission-1"));
|
||||
assert!(
|
||||
app.handle_worker_event(Event::RunEnd {
|
||||
result: RunResult::Finished,
|
||||
})
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -3501,24 +3482,6 @@ mod completion_flow_tests {
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queued_input_can_be_restored_to_composer_or_cleared() {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
insert_text(&mut app, "edit me");
|
||||
assert!(app.submit_input().is_none());
|
||||
|
||||
assert!(app.restore_next_queued_input_to_composer());
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
assert_eq!(input_text(&app), "edit me");
|
||||
|
||||
app.input.clear();
|
||||
insert_text(&mut app, "clear me");
|
||||
assert!(app.submit_input().is_none());
|
||||
assert_eq!(app.clear_queued_inputs(), 1);
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
}
|
||||
|
||||
fn insert_text(app: &mut App, text: &str) {
|
||||
for c in text.chars() {
|
||||
app.insert_char(c);
|
||||
@@ -3530,7 +3493,7 @@ mod completion_flow_tests {
|
||||
app.insert_char(c);
|
||||
}
|
||||
match app.submit_input() {
|
||||
Some(Method::Run { input }) => input,
|
||||
Some(Method::Submit { input, .. }) => input,
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
}
|
||||
}
|
||||
@@ -3675,6 +3638,7 @@ mod completion_flow_tests {
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
greeting: test_greeting(),
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
status: WorkerStatus::Running,
|
||||
@@ -3783,6 +3747,7 @@ mod completion_flow_tests {
|
||||
revision,
|
||||
status: WorkerStatus::Idle,
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
in_flight: protocol::InFlightSnapshot::default(),
|
||||
@@ -4000,6 +3965,7 @@ mod completion_flow_tests {
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
greeting: test_greeting(),
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
@@ -4051,6 +4017,7 @@ mod completion_flow_tests {
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
greeting: test_greeting(),
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
@@ -4064,6 +4031,7 @@ mod completion_flow_tests {
|
||||
},
|
||||
revision: 4,
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
status: WorkerStatus::Running,
|
||||
@@ -4222,6 +4190,7 @@ mod completion_flow_tests {
|
||||
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting,
|
||||
@@ -4437,23 +4406,23 @@ mod completion_flow_tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn input_history_records_queued_inputs_and_suppresses_consecutive_duplicates() {
|
||||
fn input_history_records_running_submits_and_suppresses_consecutive_duplicates() {
|
||||
let mut app = App::new("test".into());
|
||||
app.running = true;
|
||||
|
||||
for c in "repeat".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(app.submit_input().is_none());
|
||||
assert!(app.submit_input().is_some());
|
||||
assert_eq!(app.input_history_len(), 1);
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
|
||||
for c in "repeat".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(app.submit_input().is_none());
|
||||
assert!(app.submit_input().is_some());
|
||||
assert_eq!(app.input_history_len(), 1);
|
||||
assert_eq!(app.queued_input_count(), 2);
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
|
||||
app.insert_char(' ');
|
||||
assert!(app.submit_input().is_none());
|
||||
@@ -4481,7 +4450,7 @@ mod completion_flow_tests {
|
||||
},
|
||||
];
|
||||
app.input.replace_with_segments(&original);
|
||||
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
|
||||
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
|
||||
|
||||
assert!(app.browse_input_history_older());
|
||||
assert_eq!(app.input.submit_segments(), original);
|
||||
@@ -4493,7 +4462,7 @@ mod completion_flow_tests {
|
||||
for c in "sent".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
|
||||
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
|
||||
|
||||
for c in "draft".chars() {
|
||||
app.insert_char(c);
|
||||
@@ -4511,7 +4480,7 @@ mod completion_flow_tests {
|
||||
for c in "sent".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
|
||||
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
|
||||
|
||||
assert!(app.browse_input_history_older());
|
||||
assert!(app.input_history_is_browsing());
|
||||
@@ -4528,17 +4497,19 @@ mod completion_flow_tests {
|
||||
for c in "first".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
|
||||
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
|
||||
for c in "second".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
|
||||
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
|
||||
|
||||
assert!(app.browse_input_history_older());
|
||||
assert!(app.browse_input_history_older());
|
||||
let method = app.submit_input();
|
||||
match method {
|
||||
Some(Method::Run { input }) => assert_eq!(Segment::flatten_to_text(&input), "first"),
|
||||
Some(Method::Submit { input, .. }) => {
|
||||
assert_eq!(Segment::flatten_to_text(&input), "first")
|
||||
}
|
||||
other => panic!("expected recalled run, got {other:?}"),
|
||||
}
|
||||
assert_eq!(app.input_history_len(), 3);
|
||||
|
||||
+156
-110
@@ -270,8 +270,8 @@ impl<T: Socket> ConsoleConnection<T> {
|
||||
async fn send(&mut self, method: &Method) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut prepared = method.clone();
|
||||
let carries_attachments =
|
||||
matches!(prepared, Method::Run { .. }) && !self.pending_attachments.is_empty();
|
||||
if let Method::Run { input } = &mut prepared {
|
||||
matches!(prepared, Method::Submit { .. }) && !self.pending_attachments.is_empty();
|
||||
if let Method::Submit { input, .. } = &mut prepared {
|
||||
input.extend(
|
||||
self.pending_attachments
|
||||
.iter()
|
||||
@@ -569,6 +569,7 @@ async fn run_e2e_rewind_fixture(
|
||||
app.connected = true;
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
@@ -697,6 +698,7 @@ async fn run_e2e_rewind_fixture(
|
||||
if submitted_at.elapsed() >= apply_delay {
|
||||
app.handle_worker_event(Event::RewindApplied {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
input: vec![Segment::text("rewind-live-refresh")],
|
||||
@@ -916,7 +918,7 @@ async fn run_loop<T: Socket>(
|
||||
}
|
||||
|
||||
fn attachment_command_path(method: &Method) -> Option<PathBuf> {
|
||||
let Method::Run { input } = method else {
|
||||
let Method::Submit { input, .. } = method else {
|
||||
return None;
|
||||
};
|
||||
let [Segment::Text { content }] = input.as_slice() else {
|
||||
@@ -927,7 +929,7 @@ fn attachment_command_path(method: &Method) -> Option<PathBuf> {
|
||||
}
|
||||
|
||||
fn is_clear_attachments_command(method: &Method) -> bool {
|
||||
let Method::Run { input } = method else {
|
||||
let Method::Submit { input, .. } = method else {
|
||||
return false;
|
||||
};
|
||||
matches!(
|
||||
@@ -941,7 +943,7 @@ async fn send_console_method<T: Socket>(
|
||||
client: &mut ConsoleConnection<T>,
|
||||
method: &Method,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if matches!(method, Method::Run { .. }) && client.has_active_uploads() {
|
||||
if matches!(method, Method::Submit { .. }) && client.has_active_uploads() {
|
||||
app.restore_unsent_run(method);
|
||||
app.flash_actionbar_notice(
|
||||
"Attachment upload is still in progress; wait or use /clear-attachments.",
|
||||
@@ -953,7 +955,7 @@ async fn send_console_method<T: Socket>(
|
||||
}
|
||||
|
||||
let sends_attachments =
|
||||
matches!(method, Method::Run { .. }) && !client.pending_attachments.is_empty();
|
||||
matches!(method, Method::Submit { .. }) && !client.pending_attachments.is_empty();
|
||||
if let Err(error) = client.send(method).await {
|
||||
if sends_attachments {
|
||||
app.restore_unsent_run(method);
|
||||
@@ -1148,18 +1150,27 @@ fn handle_key(app: &mut App, key: KeyEvent) -> Option<Method> {
|
||||
app.clear_command_input();
|
||||
Some(None)
|
||||
}
|
||||
KeyCode::Char(c)
|
||||
if c.eq_ignore_ascii_case(&'d') && alt && !ctrl && !app.is_command_mode() =>
|
||||
{
|
||||
Some(
|
||||
app.next_queued_input_preview()
|
||||
.map(str::to_owned)
|
||||
.map(|submission_id| app.cancel_pending_method(submission_id)),
|
||||
)
|
||||
}
|
||||
KeyCode::Char(c)
|
||||
if c.eq_ignore_ascii_case(&'n') && alt && !ctrl && !app.is_command_mode() =>
|
||||
{
|
||||
Some(app.submit_notify_input())
|
||||
}
|
||||
KeyCode::Char(c)
|
||||
if c.eq_ignore_ascii_case(&'q') && alt && !ctrl && !app.is_command_mode() =>
|
||||
{
|
||||
if app.restore_next_queued_input_to_composer() {
|
||||
Some(app.refresh_completion())
|
||||
} else {
|
||||
Some(None)
|
||||
}
|
||||
Some(app.continue_pending_method())
|
||||
}
|
||||
KeyCode::Char(c) if c.eq_ignore_ascii_case(&'c') && alt && !ctrl => {
|
||||
app.clear_queued_inputs();
|
||||
Some(None)
|
||||
Some(Some(app.clear_pending_method()))
|
||||
}
|
||||
KeyCode::Char('c') if ctrl => Some(handle_pause_or_quit(app)),
|
||||
KeyCode::Char('x') if ctrl => Some(handle_cancel_or_shutdown(app)),
|
||||
@@ -1427,7 +1438,6 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option<Method> {
|
||||
WorkerStatus::Running | WorkerStatus::Paused
|
||||
) {
|
||||
app.shutdown_confirm = None;
|
||||
app.clear_queued_inputs();
|
||||
return Some(Method::Cancel);
|
||||
}
|
||||
if let Some(pressed_at) = app.shutdown_confirm
|
||||
@@ -1450,7 +1460,6 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option<Method> {
|
||||
/// Idle / Paused → 2-tap to quit the TUI (the Worker keeps running).
|
||||
fn handle_pause_or_quit(app: &mut App) -> Option<Method> {
|
||||
if app.worker_status == WorkerStatus::Running {
|
||||
app.clear_queued_inputs();
|
||||
return Some(Method::Pause);
|
||||
}
|
||||
if let Some(t) = app.quit_confirm
|
||||
@@ -1476,8 +1485,8 @@ mod tests {
|
||||
use crate::text_selection::{HistoryViewport, SelectionRow};
|
||||
use async_trait::async_trait;
|
||||
use protocol::{
|
||||
Event, RewindTarget, RewindTargetId, RunResult, Segment, UploadedFileAvailability,
|
||||
UploadedFileRef, WorkerStatus,
|
||||
Event, RewindTarget, RewindTargetId, Segment, UploadedFileAvailability, UploadedFileRef,
|
||||
WorkerStatus,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -1490,7 +1499,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn client_local_attachment_commands_are_typed_and_do_not_send_the_path() {
|
||||
let attach = Method::Run {
|
||||
let attach = Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: vec![Segment::text("/attach /tmp/report.md")],
|
||||
};
|
||||
assert_eq!(
|
||||
@@ -1499,7 +1509,8 @@ mod tests {
|
||||
);
|
||||
assert!(!is_clear_attachments_command(&attach));
|
||||
|
||||
let clear = Method::Run {
|
||||
let clear = Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: vec![Segment::text("/clear-attachments")],
|
||||
};
|
||||
assert!(is_clear_attachments_command(&clear));
|
||||
@@ -1605,7 +1616,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_attachment_send_failure_restores_draft_without_exiting_console() {
|
||||
async fn running_attachment_submit_failure_restores_draft_without_exiting_console() {
|
||||
let file = UploadedFileRef {
|
||||
artifact_id: "artifact-queued".into(),
|
||||
file_name: "queued.txt".into(),
|
||||
@@ -1631,13 +1642,10 @@ mod tests {
|
||||
let mut app = App::new("worker".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
app.input.insert_str("queued inspect");
|
||||
assert!(app.submit_input().is_none());
|
||||
|
||||
let method = app
|
||||
.handle_worker_event(Event::RunEnd {
|
||||
result: RunResult::Finished,
|
||||
})
|
||||
.expect("queued run must be released");
|
||||
.submit_input()
|
||||
.expect("running Submit is sent immediately");
|
||||
|
||||
send_console_method(&mut app, &mut connection, &method)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -1960,7 +1968,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn running_enter_queues_instead_of_sending_run() {
|
||||
fn running_enter_sends_submit_to_worker() {
|
||||
let mut app = App::new("agent".to_string());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
for c in "queued".chars() {
|
||||
@@ -1973,102 +1981,128 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
|
||||
assert!(matches!(
|
||||
handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)),
|
||||
Some(Method::Submit { .. })
|
||||
));
|
||||
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
assert_eq!(app.next_queued_input_preview(), Some("queued"));
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
assert_eq!(input_text(&app), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queued_input_keybindings_restore_and_clear() {
|
||||
let mut app = App::new("agent".to_string());
|
||||
fn running_alt_n_sends_explicit_notify_without_implicit_submit_conversion() {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
for c in "edit queued".chars() {
|
||||
assert!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
for character in "progress".chars() {
|
||||
app.insert_char(character);
|
||||
}
|
||||
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
|
||||
|
||||
assert!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('q'), KeyModifiers::ALT)
|
||||
)
|
||||
.is_none()
|
||||
let method = handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('n'), KeyModifiers::ALT),
|
||||
);
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
assert_eq!(input_text(&app), "edit queued");
|
||||
|
||||
app.input.clear();
|
||||
for c in "clear queued".chars() {
|
||||
assert!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
|
||||
assert!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::ALT)
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
assert!(matches!(
|
||||
method,
|
||||
Some(Method::Notify {
|
||||
ref message,
|
||||
auto_run: true,
|
||||
..
|
||||
}) if message == "progress"
|
||||
));
|
||||
assert_eq!(input_text(&app), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pause_and_cancel_clear_queued_input() {
|
||||
let mut app = App::new("agent".to_string());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
for c in "queued".chars() {
|
||||
assert!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
fn pending_queue_shortcuts_send_worker_operations() {
|
||||
let mut app = App::new("test".into());
|
||||
app.handle_worker_event(Event::PendingSubmissionsChanged {
|
||||
pending: protocol::PendingSubmissionsSnapshot {
|
||||
revision: 2,
|
||||
notification_count: 0,
|
||||
head_id: Some("submission-1".into()),
|
||||
submissions: vec![protocol::PendingSubmissionSummary {
|
||||
submission_id: "submission-1".into(),
|
||||
accepted_at_ms: 1,
|
||||
segment_count: 1,
|
||||
byte_len: 6,
|
||||
}],
|
||||
},
|
||||
});
|
||||
|
||||
let pause = handle_key(
|
||||
let continue_next = handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL),
|
||||
KeyEvent::new(KeyCode::Char('q'), KeyModifiers::ALT),
|
||||
);
|
||||
assert!(matches!(pause, Some(Method::Pause)));
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
|
||||
for c in "queued again".chars() {
|
||||
assert!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
|
||||
assert!(matches!(
|
||||
continue_next,
|
||||
Some(Method::ContinuePending {
|
||||
expected_revision: 2,
|
||||
ref expected_head_id,
|
||||
}) if expected_head_id == "submission-1"
|
||||
));
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
|
||||
let cancel = handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL),
|
||||
KeyEvent::new(KeyCode::Char('d'), KeyModifiers::ALT),
|
||||
);
|
||||
assert!(matches!(cancel, Some(Method::Cancel)));
|
||||
assert_eq!(app.queued_input_count(), 0);
|
||||
assert!(matches!(
|
||||
cancel,
|
||||
Some(Method::CancelPendingSubmission {
|
||||
expected_revision: 2,
|
||||
ref submission_id,
|
||||
}) if submission_id == "submission-1"
|
||||
));
|
||||
|
||||
let clear = handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::ALT),
|
||||
);
|
||||
assert!(matches!(
|
||||
clear,
|
||||
Some(Method::ClearPendingSubmissions {
|
||||
expected_revision: 2
|
||||
})
|
||||
));
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pause_and_cancel_preserve_authoritative_pending_queue() {
|
||||
let mut app = App::new("test".into());
|
||||
app.handle_worker_event(Event::PendingSubmissionsChanged {
|
||||
pending: protocol::PendingSubmissionsSnapshot {
|
||||
revision: 2,
|
||||
notification_count: 0,
|
||||
head_id: Some("submission-1".into()),
|
||||
submissions: vec![protocol::PendingSubmissionSummary {
|
||||
submission_id: "submission-1".into(),
|
||||
accepted_at_ms: 1,
|
||||
segment_count: 1,
|
||||
byte_len: 6,
|
||||
}],
|
||||
},
|
||||
});
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
assert!(matches!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL),
|
||||
),
|
||||
Some(Method::Pause)
|
||||
));
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
assert!(matches!(
|
||||
handle_key(
|
||||
&mut app,
|
||||
KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL),
|
||||
),
|
||||
Some(Method::Cancel)
|
||||
));
|
||||
assert_eq!(app.queued_input_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2535,13 +2569,19 @@ mod tests {
|
||||
let mut app = App::new("agent".to_string());
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
greeting: test_greeting(),
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
in_flight: Default::default(),
|
||||
internal_workers: Vec::new(),
|
||||
});
|
||||
app.handle_worker_event(Event::RewindApplied {
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
input: vec![Segment::Text {
|
||||
content: "retry this".into(),
|
||||
}],
|
||||
@@ -2562,7 +2602,10 @@ mod tests {
|
||||
let mut app = App::new("agent".to_string());
|
||||
app.handle_worker_event(Event::Snapshot {
|
||||
greeting: test_greeting(),
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
in_flight: Default::default(),
|
||||
internal_workers: Vec::new(),
|
||||
@@ -2570,7 +2613,10 @@ mod tests {
|
||||
type_keys(&mut app, "draft");
|
||||
|
||||
app.handle_worker_event(Event::RewindApplied {
|
||||
session: protocol::SessionSnapshot { entries: vec![] },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![],
|
||||
},
|
||||
input: vec![Segment::Text {
|
||||
content: "retry this".into(),
|
||||
}],
|
||||
@@ -2918,12 +2964,12 @@ mod tests {
|
||||
type_keys(&mut app, "first");
|
||||
assert!(matches!(
|
||||
handle_key(&mut app, key(KeyCode::Enter)),
|
||||
Some(Method::Run { .. })
|
||||
Some(Method::Submit { .. })
|
||||
));
|
||||
type_keys(&mut app, "second");
|
||||
assert!(matches!(
|
||||
handle_key(&mut app, key(KeyCode::Enter)),
|
||||
Some(Method::Run { .. })
|
||||
Some(Method::Submit { .. })
|
||||
));
|
||||
|
||||
assert_eq!(input_text(&app), "");
|
||||
@@ -2954,7 +3000,7 @@ mod tests {
|
||||
type_keys(&mut app, "sent");
|
||||
assert!(matches!(
|
||||
handle_key(&mut app, key(KeyCode::Enter)),
|
||||
Some(Method::Run { .. })
|
||||
Some(Method::Submit { .. })
|
||||
));
|
||||
type_keys(&mut app, "draft\nbody");
|
||||
app.move_cursor_start();
|
||||
|
||||
+22
-14
@@ -1880,7 +1880,7 @@ fn actionbar_left_item(app: &App, now: Instant) -> Option<(String, Style)> {
|
||||
}
|
||||
if app.queued_input_count() > 0 {
|
||||
return Some((
|
||||
"Alt-q edit queued Alt-c clear queued".to_string(),
|
||||
"Alt-n notify Alt-q continue Alt-d cancel queued Alt-c clear queued".to_string(),
|
||||
Style::default().fg(Color::DarkGray),
|
||||
));
|
||||
}
|
||||
@@ -2136,9 +2136,25 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::app::{ActionbarNoticeLevel, ActionbarNoticeSource, App};
|
||||
use crate::block::{ToolCallBlock, ToolCallState};
|
||||
use protocol::WorkerStatus;
|
||||
use protocol::Event;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
fn set_pending_submission(app: &mut App, id: &str) {
|
||||
app.handle_worker_event(Event::PendingSubmissionsChanged {
|
||||
pending: protocol::PendingSubmissionsSnapshot {
|
||||
revision: 1,
|
||||
notification_count: 0,
|
||||
head_id: Some(id.into()),
|
||||
submissions: vec![protocol::PendingSubmissionSummary {
|
||||
submission_id: id.into(),
|
||||
accepted_at_ms: 1,
|
||||
segment_count: 1,
|
||||
byte_len: 1,
|
||||
}],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn run_status_line_matches_console_metrics_and_spinner_frame() {
|
||||
let now = Instant::now();
|
||||
@@ -2251,15 +2267,11 @@ mod tests {
|
||||
#[test]
|
||||
fn queue_status_text_includes_count_and_preview() {
|
||||
let mut app = App::new("test".into());
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
for c in "queued preview".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(app.submit_input().is_none());
|
||||
set_pending_submission(&mut app, "submission-1");
|
||||
|
||||
assert_eq!(
|
||||
queue_status_text(&app),
|
||||
Some("queued: 1 — queued preview".to_string())
|
||||
Some("queued: 1 — submission-1".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2289,14 +2301,10 @@ mod tests {
|
||||
Some("Worker keeps running. Press Ctrl-C again to exit TUI.".into())
|
||||
);
|
||||
|
||||
app.set_worker_status(WorkerStatus::Running);
|
||||
for c in "queued turn".chars() {
|
||||
app.insert_char(c);
|
||||
}
|
||||
assert!(app.submit_input().is_none());
|
||||
set_pending_submission(&mut app, "submission-1");
|
||||
assert_eq!(
|
||||
actionbar_left_item(&app, now).map(|(text, _)| text),
|
||||
Some("Alt-q edit queued Alt-c clear queued".into())
|
||||
Some("Alt-n notify Alt-q continue Alt-d cancel queued Alt-c clear queued".into())
|
||||
);
|
||||
|
||||
app.enter_command_mode();
|
||||
|
||||
@@ -41,14 +41,12 @@ pub enum WorkerExecutionOperation {
|
||||
Cancel,
|
||||
}
|
||||
|
||||
/// Evidence that a user input reached the durable Worker session boundary.
|
||||
///
|
||||
/// This is intentionally distinct from accepting a method on the Worker's
|
||||
/// in-memory channel. For Flow submissions, the committed UserInput entry also
|
||||
/// carries the initial Flow runtime-state extension.
|
||||
/// Evidence that a Submit request reached the durable Worker session boundary.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct WorkerInputCommitAck {
|
||||
pub struct WorkerSubmissionAck {
|
||||
pub submission_request_id: String,
|
||||
pub submission_id: String,
|
||||
pub disposition: protocol::SubmissionDisposition,
|
||||
}
|
||||
|
||||
/// Typed execution result class. Results are transient operation outcomes and
|
||||
@@ -61,7 +59,7 @@ pub struct WorkerExecutionResult {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub message: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub input_commit: Option<WorkerInputCommitAck>,
|
||||
pub submission: Option<WorkerSubmissionAck>,
|
||||
}
|
||||
|
||||
/// Backend result class for a Worker execution operation.
|
||||
@@ -85,22 +83,26 @@ impl WorkerExecutionResult {
|
||||
outcome: WorkerExecutionOutcome::Accepted,
|
||||
run_state,
|
||||
message: None,
|
||||
input_commit: None,
|
||||
submission: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn accepted_input_committed(
|
||||
pub fn accepted_submission(
|
||||
operation: WorkerExecutionOperation,
|
||||
run_state: WorkerExecutionRunState,
|
||||
submission_request_id: impl Into<String>,
|
||||
submission_id: impl Into<String>,
|
||||
disposition: protocol::SubmissionDisposition,
|
||||
) -> Self {
|
||||
Self {
|
||||
operation,
|
||||
outcome: WorkerExecutionOutcome::Accepted,
|
||||
run_state,
|
||||
message: None,
|
||||
input_commit: Some(WorkerInputCommitAck {
|
||||
submission: Some(WorkerSubmissionAck {
|
||||
submission_request_id: submission_request_id.into(),
|
||||
submission_id: submission_id.into(),
|
||||
disposition,
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -111,7 +113,7 @@ impl WorkerExecutionResult {
|
||||
outcome: WorkerExecutionOutcome::Busy,
|
||||
run_state: WorkerExecutionRunState::Busy,
|
||||
message: Some(message.into()),
|
||||
input_commit: None,
|
||||
submission: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,7 +123,7 @@ impl WorkerExecutionResult {
|
||||
outcome: WorkerExecutionOutcome::Rejected,
|
||||
run_state: WorkerExecutionRunState::Stopped,
|
||||
message: Some(message.into()),
|
||||
input_commit: None,
|
||||
submission: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -131,7 +133,7 @@ impl WorkerExecutionResult {
|
||||
outcome: WorkerExecutionOutcome::Errored,
|
||||
run_state: WorkerExecutionRunState::Errored,
|
||||
message: Some(message.into()),
|
||||
input_commit: None,
|
||||
submission: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,7 +143,7 @@ impl WorkerExecutionResult {
|
||||
outcome: WorkerExecutionOutcome::Unsupported,
|
||||
run_state: WorkerExecutionRunState::Stopped,
|
||||
message: Some(message.into()),
|
||||
input_commit: None,
|
||||
submission: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -618,14 +620,17 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn input_commit_ack_survives_json_round_trip() {
|
||||
let result = WorkerExecutionResult::accepted_input_committed(
|
||||
fn submission_ack_survives_json_round_trip() {
|
||||
let result = WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Busy,
|
||||
"request-1",
|
||||
"submission-1",
|
||||
protocol::SubmissionDisposition::Started,
|
||||
);
|
||||
|
||||
let json = serde_json::to_string(&result).unwrap();
|
||||
assert!(json.contains("\"submission_request_id\":\"request-1\""));
|
||||
assert!(json.contains("\"submission_id\":\"submission-1\""));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<WorkerExecutionResult>(&json).unwrap(),
|
||||
|
||||
@@ -1240,10 +1240,12 @@ async fn worker_protocol_ws(
|
||||
auth: Option<Extension<RuntimeAuthContext>>,
|
||||
Path(worker_id): Path<String>,
|
||||
Query(query): Query<RuntimeWorkerEventsWsQuery>,
|
||||
headers: HeaderMap,
|
||||
ws: WebSocketUpgrade,
|
||||
) -> Result<Response, RuntimeHttpRestError> {
|
||||
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
|
||||
let scope = auth_workspace_scope(&state, auth.as_ref())?;
|
||||
let input_source = authenticated_protocol_input_source(&headers)?;
|
||||
match scope.as_ref() {
|
||||
Some(scope) => state
|
||||
.runtime
|
||||
@@ -1254,17 +1256,89 @@ async fn worker_protocol_ws(
|
||||
.map_err(RuntimeHttpRestError::runtime)?;
|
||||
Ok(ws
|
||||
.on_upgrade(move |socket| {
|
||||
worker_protocol_ws_session(state.runtime, scope, worker_ref, query, socket)
|
||||
worker_protocol_ws_session(
|
||||
state.runtime,
|
||||
scope,
|
||||
worker_ref,
|
||||
query,
|
||||
input_source,
|
||||
socket,
|
||||
)
|
||||
})
|
||||
.into_response())
|
||||
}
|
||||
|
||||
#[cfg(feature = "ws-server")]
|
||||
fn authenticated_protocol_input_source(
|
||||
headers: &HeaderMap,
|
||||
) -> Result<Option<protocol::AuthenticatedInputSource>, RuntimeHttpRestError> {
|
||||
let Some(value) = headers.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let account_id = value.to_str().map_err(|_| {
|
||||
RuntimeHttpRestError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"authenticated_input_source_invalid",
|
||||
"authenticated Worker input source is invalid",
|
||||
)
|
||||
})?;
|
||||
if account_id.trim().is_empty() || account_id.len() > 128 {
|
||||
return Err(RuntimeHttpRestError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"authenticated_input_source_invalid",
|
||||
"authenticated Worker input source is invalid",
|
||||
));
|
||||
}
|
||||
Ok(Some(protocol::AuthenticatedInputSource::Account {
|
||||
account_id: account_id.to_owned(),
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(feature = "ws-server")]
|
||||
fn authorize_runtime_protocol_method(
|
||||
method: protocol::Method,
|
||||
transport_source: Option<&protocol::AuthenticatedInputSource>,
|
||||
) -> protocol::Method {
|
||||
match method {
|
||||
protocol::Method::SubmitTracked {
|
||||
submission_request_id,
|
||||
input,
|
||||
..
|
||||
} => protocol::Method::SubmitTracked {
|
||||
source: transport_source.cloned().unwrap_or_else(|| {
|
||||
protocol::AuthenticatedInputSource::Backend {
|
||||
operation_id: submission_request_id.clone(),
|
||||
}
|
||||
}),
|
||||
submission_request_id,
|
||||
input,
|
||||
},
|
||||
protocol::Method::NotifyTracked {
|
||||
notification_request_id,
|
||||
message,
|
||||
auto_run,
|
||||
..
|
||||
} => protocol::Method::NotifyTracked {
|
||||
source: transport_source.cloned().unwrap_or_else(|| {
|
||||
protocol::AuthenticatedInputSource::Backend {
|
||||
operation_id: notification_request_id.clone(),
|
||||
}
|
||||
}),
|
||||
notification_request_id,
|
||||
message,
|
||||
auto_run,
|
||||
},
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "ws-server")]
|
||||
async fn worker_protocol_ws_session(
|
||||
runtime: Runtime,
|
||||
scope: Option<RuntimeWorkspaceScope>,
|
||||
worker_ref: WorkerRef,
|
||||
query: RuntimeWorkerEventsWsQuery,
|
||||
input_source: Option<protocol::AuthenticatedInputSource>,
|
||||
mut socket: WebSocket,
|
||||
) {
|
||||
let mut cursor = match query.cursor.as_deref() {
|
||||
@@ -1347,6 +1421,8 @@ async fn worker_protocol_ws_session(
|
||||
match inbound {
|
||||
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
|
||||
Ok(method) => {
|
||||
let method =
|
||||
authorize_runtime_protocol_method(method, input_source.as_ref());
|
||||
let result = match scope.as_ref() {
|
||||
Some(scope) => {
|
||||
runtime.send_protocol_method_scoped(scope, &worker_ref, method)
|
||||
@@ -2219,6 +2295,63 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_protocol_replaces_serialized_tracked_source() {
|
||||
let wire = serde_json::to_string(&protocol::Method::SubmitTracked {
|
||||
submission_request_id: "request-1".into(),
|
||||
input: vec![protocol::Segment::text("hello")],
|
||||
source: protocol::AuthenticatedInputSource::Account {
|
||||
account_id: "forged".into(),
|
||||
},
|
||||
})
|
||||
.unwrap();
|
||||
let decoded: protocol::Method = serde_json::from_str(&wire).unwrap();
|
||||
assert!(matches!(
|
||||
decoded,
|
||||
protocol::Method::SubmitTracked {
|
||||
source: protocol::AuthenticatedInputSource::UntrustedWire,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert!(matches!(
|
||||
authorize_runtime_protocol_method(decoded, None),
|
||||
protocol::Method::SubmitTracked {
|
||||
source: protocol::AuthenticatedInputSource::Backend { operation_id },
|
||||
..
|
||||
} if operation_id == "request-1"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_protocol_uses_transport_authenticated_account_source() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
|
||||
"account-1".parse().unwrap(),
|
||||
);
|
||||
let source = authenticated_protocol_input_source(&headers)
|
||||
.unwrap()
|
||||
.expect("account source header must resolve");
|
||||
let wire = serde_json::to_string(&protocol::Method::NotifyTracked {
|
||||
notification_request_id: "notification-1".into(),
|
||||
message: "hello".into(),
|
||||
auto_run: true,
|
||||
source: protocol::AuthenticatedInputSource::Account {
|
||||
account_id: "forged".into(),
|
||||
},
|
||||
})
|
||||
.unwrap();
|
||||
let decoded: protocol::Method = serde_json::from_str(&wire).unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
authorize_runtime_protocol_method(decoded, Some(&source)),
|
||||
protocol::Method::NotifyTracked {
|
||||
source: protocol::AuthenticatedInputSource::Account { account_id },
|
||||
..
|
||||
} if account_id == "account-1"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attachment_routes_require_worker_input_permission() {
|
||||
assert_eq!(
|
||||
@@ -2870,11 +3003,13 @@ mod tests {
|
||||
_handle: &WorkerExecutionHandle,
|
||||
input: WorkerInput,
|
||||
) -> WorkerExecutionResult {
|
||||
if let Some(submission_id) = input.submission_id {
|
||||
WorkerExecutionResult::accepted_input_committed(
|
||||
if let Some(submission_id) = input.submission_request_id {
|
||||
WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Idle,
|
||||
submission_id.clone(),
|
||||
submission_id,
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
} else {
|
||||
WorkerExecutionResult::accepted(
|
||||
@@ -3194,11 +3329,13 @@ mod ws_tests {
|
||||
_handle: &WorkerExecutionHandle,
|
||||
input: WorkerInput,
|
||||
) -> WorkerExecutionResult {
|
||||
if let Some(submission_id) = input.submission_id {
|
||||
WorkerExecutionResult::accepted_input_committed(
|
||||
if let Some(submission_id) = input.submission_request_id {
|
||||
WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Idle,
|
||||
submission_id.clone(),
|
||||
submission_id,
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
} else {
|
||||
WorkerExecutionResult::accepted(
|
||||
|
||||
@@ -25,10 +25,10 @@ impl WorkerInputKind {
|
||||
pub struct WorkerInput {
|
||||
pub kind: WorkerInputKind,
|
||||
pub content: String,
|
||||
/// Runtime-generated correlation id. This is never accepted from public
|
||||
/// JSON input and is consumed only by the execution backend.
|
||||
#[serde(skip)]
|
||||
pub submission_id: Option<String>,
|
||||
/// Authenticated client-generated idempotency key. Runtime generates one
|
||||
/// only for trusted internal callers that omit it.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub submission_request_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub segments: Option<Vec<Segment>>,
|
||||
}
|
||||
@@ -38,7 +38,7 @@ impl WorkerInput {
|
||||
Self {
|
||||
kind: WorkerInputKind::User,
|
||||
content: content.into(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: None,
|
||||
}
|
||||
}
|
||||
@@ -47,7 +47,7 @@ impl WorkerInput {
|
||||
Self {
|
||||
kind: WorkerInputKind::Notify,
|
||||
content: content.into(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: None,
|
||||
}
|
||||
}
|
||||
@@ -57,6 +57,21 @@ impl WorkerInput {
|
||||
mod tests {
|
||||
use super::WorkerInput;
|
||||
|
||||
#[test]
|
||||
fn submission_request_id_round_trips_for_authenticated_client_retry() {
|
||||
let input: WorkerInput = serde_json::from_value(serde_json::json!({
|
||||
"kind": "user",
|
||||
"content": "message",
|
||||
"submission_request_id": "request-1"
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(input.submission_request_id.as_deref(), Some("request-1"));
|
||||
assert_eq!(
|
||||
serde_json::to_value(input).unwrap()["submission_request_id"],
|
||||
"request-1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notify_is_an_operation_and_legacy_system_kind_is_rejected() {
|
||||
assert_eq!(
|
||||
@@ -78,4 +93,7 @@ mod tests {
|
||||
pub struct WorkerInteractionAck {
|
||||
pub worker_ref: WorkerRef,
|
||||
pub status: WorkerStatus,
|
||||
/// Present for User Submit and absent for non-Submit interactions.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub submission: Option<crate::execution::WorkerSubmissionAck>,
|
||||
}
|
||||
|
||||
@@ -748,8 +748,12 @@ impl Runtime {
|
||||
let state = self.lock()?;
|
||||
state.worker(&worker_ref)?.request.initial_input.clone()
|
||||
} {
|
||||
let expected_submission_id = Uuid::now_v7().to_string();
|
||||
initial_input.submission_id = Some(expected_submission_id.clone());
|
||||
let expected_submission_id = initial_input
|
||||
.submission_request_id
|
||||
.clone()
|
||||
.filter(|request_id| !request_id.trim().is_empty())
|
||||
.unwrap_or_else(|| Uuid::now_v7().to_string());
|
||||
initial_input.submission_request_id = Some(expected_submission_id.clone());
|
||||
let dispatch_result = backend.dispatch_input(&handle, initial_input.clone());
|
||||
if !dispatch_result.is_accepted() {
|
||||
let _ = backend.stop_worker(&handle);
|
||||
@@ -763,9 +767,9 @@ impl Runtime {
|
||||
});
|
||||
}
|
||||
let has_commit_ack = dispatch_result
|
||||
.input_commit
|
||||
.submission
|
||||
.as_ref()
|
||||
.is_some_and(|ack| ack.submission_id == expected_submission_id);
|
||||
.is_some_and(|ack| ack.submission_request_id == expected_submission_id);
|
||||
if !has_commit_ack {
|
||||
let _ = backend.stop_worker(&handle);
|
||||
self.rollback_failed_create(&worker_ref)?;
|
||||
@@ -1146,13 +1150,18 @@ impl Runtime {
|
||||
mut input: WorkerInput,
|
||||
) -> Result<WorkerInteractionAck, RuntimeError> {
|
||||
validate_worker_input(&input)?;
|
||||
let expected_submission_id = if input.kind == WorkerInputKind::User {
|
||||
let submission_id = Uuid::now_v7().to_string();
|
||||
input.submission_id = Some(submission_id.clone());
|
||||
Some(submission_id)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let expected_submission_id =
|
||||
if matches!(input.kind, WorkerInputKind::User | WorkerInputKind::Notify) {
|
||||
let submission_id = input
|
||||
.submission_request_id
|
||||
.clone()
|
||||
.filter(|request_id| !request_id.trim().is_empty())
|
||||
.unwrap_or_else(|| Uuid::now_v7().to_string());
|
||||
input.submission_request_id = Some(submission_id.clone());
|
||||
Some(submission_id)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
self.ensure_worker_execution(worker_ref)?;
|
||||
let (backend, handle) = {
|
||||
let state = self.lock()?;
|
||||
@@ -1191,13 +1200,13 @@ impl Runtime {
|
||||
}
|
||||
if let Some(expected_submission_id) = expected_submission_id
|
||||
&& dispatch_result
|
||||
.input_commit
|
||||
.submission
|
||||
.as_ref()
|
||||
.is_none_or(|ack| ack.submission_id != expected_submission_id)
|
||||
.is_none_or(|ack| ack.submission_request_id != expected_submission_id)
|
||||
{
|
||||
let result = WorkerExecutionResult::rejected(
|
||||
WorkerExecutionOperation::Input,
|
||||
"execution backend did not acknowledge the committed Runtime submission id",
|
||||
"execution backend did not acknowledge the committed Runtime submission request id",
|
||||
);
|
||||
self.record_execution_result(worker_ref, result.clone())?;
|
||||
return Err(RuntimeError::WorkerExecutionRejected {
|
||||
@@ -1209,6 +1218,7 @@ impl Runtime {
|
||||
});
|
||||
}
|
||||
|
||||
let submission = dispatch_result.submission.clone();
|
||||
let mut state = self.lock()?;
|
||||
state.ensure_running()?;
|
||||
let worker = state.worker_mut(worker_ref)?;
|
||||
@@ -1225,6 +1235,7 @@ impl Runtime {
|
||||
Ok(WorkerInteractionAck {
|
||||
worker_ref: worker_ref.clone(),
|
||||
status,
|
||||
submission,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1706,6 +1717,7 @@ impl Runtime {
|
||||
}
|
||||
Ok(protocol::Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -3250,17 +3262,9 @@ fn validate_worker_input(input: &WorkerInput) -> Result<(), RuntimeError> {
|
||||
#[cfg(feature = "ws-server")]
|
||||
fn input_protocol_event(input: &WorkerInput) -> Option<protocol::Event> {
|
||||
match input.kind {
|
||||
WorkerInputKind::User => Some(protocol::Event::UserMessage {
|
||||
segments: input.segments.clone().unwrap_or_else(|| {
|
||||
vec![protocol::Segment::Text {
|
||||
content: input.content.clone(),
|
||||
}]
|
||||
}),
|
||||
}),
|
||||
// The committed `SystemItem::Notification` is the sole agent-visible
|
||||
// and Console-visible authority for Notify. A synthetic observation
|
||||
// here would display the same notification twice.
|
||||
WorkerInputKind::Notify => None,
|
||||
// Submit is projected only after the Worker commits UserInput. Queued
|
||||
// payloads must never become model- or client-visible history early.
|
||||
WorkerInputKind::User | WorkerInputKind::Notify => None,
|
||||
WorkerInputKind::Compact
|
||||
| WorkerInputKind::ListRewindTargets
|
||||
| WorkerInputKind::RegisterPeer => Some(protocol::Event::SystemItem {
|
||||
@@ -3435,6 +3439,7 @@ mod tests {
|
||||
);
|
||||
let snapshot = protocol::Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -3475,7 +3480,7 @@ mod tests {
|
||||
let input = WorkerInput {
|
||||
kind: WorkerInputKind::User,
|
||||
content: String::new(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: Some(vec![protocol::Segment::Flow {
|
||||
selector: "builtin:coder-review".to_string(),
|
||||
}]),
|
||||
@@ -3488,7 +3493,7 @@ mod tests {
|
||||
let input = WorkerInput {
|
||||
kind: WorkerInputKind::User,
|
||||
content: String::new(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: Some(Vec::new()),
|
||||
};
|
||||
assert!(matches!(
|
||||
@@ -3503,7 +3508,7 @@ mod tests {
|
||||
request.initial_input = Some(WorkerInput {
|
||||
kind: WorkerInputKind::User,
|
||||
content: String::new(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: Some(vec![protocol::Segment::Flow {
|
||||
selector: "builtin:coder-review".to_string(),
|
||||
}]),
|
||||
@@ -4005,7 +4010,7 @@ mod tests {
|
||||
_handle: &WorkerExecutionHandle,
|
||||
input: WorkerInput,
|
||||
) -> WorkerExecutionResult {
|
||||
let submission_id = input.submission_id.clone();
|
||||
let submission_id = input.submission_request_id.clone();
|
||||
self.dispatched_inputs.lock().unwrap().push(input);
|
||||
let mut result = self
|
||||
.dispatch_result
|
||||
@@ -4013,19 +4018,21 @@ mod tests {
|
||||
.unwrap()
|
||||
.clone()
|
||||
.unwrap_or_else(|| {
|
||||
WorkerExecutionResult::accepted_input_committed(
|
||||
WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Idle,
|
||||
"request-test",
|
||||
"test-submission",
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
});
|
||||
if !self
|
||||
.preserve_commit_ack_submission_id
|
||||
.load(Ordering::SeqCst)
|
||||
&& let (Some(ack), Some(submission_id)) =
|
||||
(result.input_commit.as_mut(), submission_id)
|
||||
(result.submission.as_mut(), submission_id)
|
||||
{
|
||||
ack.submission_id = submission_id;
|
||||
ack.submission_request_id = submission_id;
|
||||
}
|
||||
result
|
||||
}
|
||||
@@ -4717,10 +4724,12 @@ mod tests {
|
||||
#[test]
|
||||
fn create_worker_uses_committed_input_ack_run_state() {
|
||||
let (runtime, backend) = runtime_and_backend();
|
||||
backend.set_dispatch_result(WorkerExecutionResult::accepted_input_committed(
|
||||
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Idle,
|
||||
"request-test",
|
||||
"test-submission",
|
||||
protocol::SubmissionDisposition::Started,
|
||||
));
|
||||
let mut request = task_request("committed initial input is already idle");
|
||||
request.initial_input = Some(WorkerInput::user("start the ticket"));
|
||||
@@ -4731,13 +4740,15 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_worker_rejects_mismatched_input_commit_acknowledgement() {
|
||||
fn create_worker_rejects_mismatched_submission_acknowledgement() {
|
||||
let (runtime, backend) = runtime_and_backend();
|
||||
backend.preserve_commit_ack_submission_id();
|
||||
backend.set_dispatch_result(WorkerExecutionResult::accepted_input_committed(
|
||||
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Busy,
|
||||
"request-test",
|
||||
"forged-submission",
|
||||
protocol::SubmissionDisposition::Started,
|
||||
));
|
||||
let mut request = task_request("mismatched initial input commit ack");
|
||||
request.initial_input = Some(WorkerInput::user("start the ticket"));
|
||||
@@ -4866,6 +4877,7 @@ mod tests {
|
||||
&detail.worker_ref,
|
||||
protocol::Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: vec![protocol::SessionSnapshotEntry {
|
||||
entry_id: "restored-log-entry".to_owned(),
|
||||
timestamp: 1,
|
||||
@@ -4937,10 +4949,14 @@ mod tests {
|
||||
_handle: &WorkerExecutionHandle,
|
||||
input: WorkerInput,
|
||||
) -> WorkerExecutionResult {
|
||||
WorkerExecutionResult::accepted_input_committed(
|
||||
WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Idle,
|
||||
input.submission_id.expect("Runtime submission id"),
|
||||
"request-test",
|
||||
input
|
||||
.submission_request_id
|
||||
.expect("Runtime submission request id"),
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -5030,7 +5046,7 @@ mod tests {
|
||||
request.initial_input = Some(WorkerInput {
|
||||
kind: WorkerInputKind::User,
|
||||
content: String::new(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: Some(vec![
|
||||
protocol::Segment::Flow {
|
||||
selector: "builtin:coder-review".to_string(),
|
||||
@@ -5068,7 +5084,7 @@ mod tests {
|
||||
let input = WorkerInput {
|
||||
kind: WorkerInputKind::User,
|
||||
content: String::new(),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: Some(vec![protocol::Segment::Flow {
|
||||
selector: "builtin:coder-review".to_string(),
|
||||
}]),
|
||||
@@ -5083,11 +5099,11 @@ mod tests {
|
||||
assert_eq!(dispatched[0].kind, input.kind);
|
||||
assert_eq!(dispatched[0].content, input.content);
|
||||
assert_eq!(dispatched[0].segments, input.segments);
|
||||
let submission_id = dispatched[0]
|
||||
.submission_id
|
||||
let submission_request_id = dispatched[0]
|
||||
.submission_request_id
|
||||
.as_deref()
|
||||
.expect("Runtime submission id");
|
||||
Uuid::parse_str(submission_id).expect("submission id UUID");
|
||||
.expect("Runtime submission request id");
|
||||
Uuid::parse_str(submission_request_id).expect("submission request id UUID");
|
||||
}
|
||||
|
||||
#[cfg(feature = "ws-server")]
|
||||
@@ -5113,11 +5129,7 @@ mod tests {
|
||||
let observations = runtime
|
||||
.read_worker_observation_events(&detail.worker_ref, WorkerObservationCursor::zero())
|
||||
.unwrap();
|
||||
assert_eq!(observations.len(), 1);
|
||||
assert!(matches!(
|
||||
observations[0].payload,
|
||||
protocol::Event::UserMessage { .. }
|
||||
));
|
||||
assert!(observations.is_empty());
|
||||
|
||||
runtime
|
||||
.observe_worker_event(
|
||||
@@ -5135,8 +5147,8 @@ mod tests {
|
||||
let observations = runtime
|
||||
.read_worker_observation_events(&detail.worker_ref, WorkerObservationCursor::zero())
|
||||
.unwrap();
|
||||
assert_eq!(observations.len(), 2);
|
||||
let protocol::Event::SystemItem { item } = &observations[1].payload else {
|
||||
assert_eq!(observations.len(), 1);
|
||||
let protocol::Event::SystemItem { item } = &observations[0].payload else {
|
||||
panic!("committed notification observation must be a system item");
|
||||
};
|
||||
assert_eq!(item["kind"], "notification");
|
||||
|
||||
@@ -39,7 +39,7 @@ use crate::working_directory::{
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use protocol::{ErrorCode, Event, Method, Segment, WorkerStatus};
|
||||
use session_store::{CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore};
|
||||
use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore};
|
||||
#[cfg(test)]
|
||||
use session_store::{FsStore, FsWorkerStore};
|
||||
use tokio::runtime::Runtime;
|
||||
@@ -57,11 +57,10 @@ use worker::feature::builtin::{
|
||||
#[cfg(feature = "ws-server")]
|
||||
use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session};
|
||||
use worker::{
|
||||
PreparedWorker, PromptCatalogSource, SegmentLogSink, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN,
|
||||
Worker, WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout,
|
||||
WorkerControllerTransport, WorkerError, WorkerFilesystemAuthority, WorkerHandle,
|
||||
WorkerSharedState, WorkerWorkspaceContext, WorkspaceClient, WorkspaceId,
|
||||
bash_output_dir_for_worker_id,
|
||||
PreparedWorker, PromptCatalogSource, SegmentLogSink, Worker, WorkerBootstrap,
|
||||
WorkerBootstrapError, WorkerBootstrapLayout, WorkerControllerTransport, WorkerError,
|
||||
WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState, WorkerWorkspaceContext,
|
||||
WorkspaceClient, WorkspaceId, bash_output_dir_for_worker_id,
|
||||
};
|
||||
|
||||
const DEFAULT_BACKEND_ID: &str = "worker-crate";
|
||||
@@ -70,17 +69,6 @@ const RUNTIME_TASK_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
// returns a typed execution error instead of leaving the outer waiter to time out.
|
||||
const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9);
|
||||
|
||||
fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool {
|
||||
let extensions = match entry {
|
||||
LogEntry::AnnotatedUserInput { extensions, .. } => extensions,
|
||||
_ => return false,
|
||||
};
|
||||
extensions.iter().any(|extension| {
|
||||
extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN
|
||||
&& extension.payload["submission_id"].as_str() == Some(submission_id)
|
||||
})
|
||||
}
|
||||
|
||||
pub struct RuntimeWorkerController {
|
||||
pub handle: WorkerHandle,
|
||||
pub shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
|
||||
@@ -1342,126 +1330,75 @@ where
|
||||
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
|
||||
}
|
||||
|
||||
fn send_user_input_and_wait_for_commit(
|
||||
fn send_submit_and_wait_for_acceptance(
|
||||
&self,
|
||||
operation: WorkerExecutionOperation,
|
||||
worker: WorkerHandle,
|
||||
method: Method,
|
||||
submission_id: String,
|
||||
submission_request_id: String,
|
||||
accepted_run_state: WorkerExecutionRunState,
|
||||
) -> WorkerExecutionResult {
|
||||
let acknowledged_submission_id = submission_id.clone();
|
||||
let request_id = submission_request_id.clone();
|
||||
self.run_on_adapter_runtime(async move {
|
||||
// Subscribe before enqueueing the input so the acknowledgement cannot
|
||||
// race with a fast Worker commit. The opaque submission id is stored in
|
||||
// the same UserInput entry as the transformed Flow input and its state.
|
||||
let (_, mut committed_entries) = worker.sink.subscribe_with_snapshot();
|
||||
let committed_probe = worker.clone();
|
||||
// Subscribe before enqueueing so a fast durable acceptance cannot
|
||||
// race the Runtime acknowledgement.
|
||||
let mut events = worker.subscribe();
|
||||
worker
|
||||
.send(method)
|
||||
.await
|
||||
.map_err(|err| format!("failed to send Worker method: {err}"))?;
|
||||
|
||||
let timeout_probe = committed_probe.clone();
|
||||
let timeout_submission_id = submission_id.clone();
|
||||
let acknowledgement = tokio::time::timeout(USER_INPUT_COMMIT_TIMEOUT, async move {
|
||||
let input_was_committed = || {
|
||||
committed_probe
|
||||
.committed_entries()
|
||||
.iter()
|
||||
.any(|entry| user_input_has_submission(entry, &submission_id))
|
||||
};
|
||||
tokio::time::timeout(USER_INPUT_COMMIT_TIMEOUT, async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
entry = committed_entries.recv() => {
|
||||
match entry {
|
||||
Ok(entry) if user_input_has_submission(&entry, &submission_id) => {
|
||||
return Ok(());
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
|
||||
if input_was_committed() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(format!(
|
||||
"worker input commit acknowledgement lagged by {skipped} entry event(s)"
|
||||
));
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
|
||||
if input_was_committed() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(
|
||||
"worker entry stream closed before user input was committed"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
match events.recv().await {
|
||||
Ok(Event::SubmissionAccepted {
|
||||
submission_request_id,
|
||||
submission_id,
|
||||
disposition,
|
||||
}) if submission_request_id == request_id => {
|
||||
return Ok((submission_id, disposition));
|
||||
}
|
||||
event = events.recv() => {
|
||||
match event {
|
||||
Ok(Event::Error { message, .. }) => {
|
||||
if input_was_committed() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(format!(
|
||||
"worker rejected user input before session commit: {message}"
|
||||
));
|
||||
}
|
||||
Ok(Event::Shutdown) => {
|
||||
if input_was_committed() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(
|
||||
"worker shut down before user input was committed".to_string()
|
||||
);
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
|
||||
if input_was_committed() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(format!(
|
||||
"worker input commit acknowledgement lagged by {skipped} protocol event(s)"
|
||||
));
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
|
||||
if input_was_committed() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(
|
||||
"worker event stream closed before user input was committed"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(Event::SubmissionRejected {
|
||||
submission_request_id,
|
||||
message,
|
||||
}) if submission_request_id == request_id => {
|
||||
return Err(format!("worker rejected Submit: {message}"));
|
||||
}
|
||||
Ok(Event::Error { message, .. }) => {
|
||||
return Err(format!(
|
||||
"worker rejected Submit before durable acceptance: {message}"
|
||||
));
|
||||
}
|
||||
Ok(Event::Shutdown) => {
|
||||
return Err(
|
||||
"worker shut down before Submit was durably accepted".to_string()
|
||||
);
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
|
||||
return Err(format!(
|
||||
"worker Submit acknowledgement lagged by {skipped} protocol event(s)"
|
||||
));
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
|
||||
return Err(
|
||||
"worker event stream closed before Submit was durably accepted"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
match acknowledgement {
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
if timeout_probe
|
||||
.committed_entries()
|
||||
.iter()
|
||||
.any(|entry| user_input_has_submission(entry, &timeout_submission_id))
|
||||
{
|
||||
Ok(())
|
||||
} else {
|
||||
Err("timed out waiting for worker user input commit".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
.await
|
||||
.map_err(|_| "timed out waiting for durable Worker Submit acceptance".to_string())?
|
||||
})
|
||||
.map(|_| {
|
||||
WorkerExecutionResult::accepted_input_committed(
|
||||
.map(|(submission_id, disposition)| {
|
||||
WorkerExecutionResult::accepted_submission(
|
||||
operation,
|
||||
accepted_run_state,
|
||||
acknowledged_submission_id,
|
||||
submission_request_id,
|
||||
submission_id,
|
||||
disposition,
|
||||
)
|
||||
})
|
||||
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
|
||||
@@ -1582,9 +1519,10 @@ impl<F> Drop for WorkerRuntimeExecutionBackend<F> {
|
||||
fn method_starts_turn(method: &Method) -> bool {
|
||||
matches!(
|
||||
method,
|
||||
Method::Run { .. }
|
||||
| Method::RunTracked { .. }
|
||||
Method::Submit { .. }
|
||||
| Method::SubmitTracked { .. }
|
||||
| Method::Notify { auto_run: true, .. }
|
||||
| Method::NotifyTracked { auto_run: true, .. }
|
||||
| Method::Resume
|
||||
| Method::Compact
|
||||
)
|
||||
@@ -1609,9 +1547,10 @@ fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExec
|
||||
|
||||
fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState {
|
||||
match method {
|
||||
Method::Run { .. }
|
||||
| Method::RunTracked { .. }
|
||||
Method::Submit { .. }
|
||||
| Method::SubmitTracked { .. }
|
||||
| Method::Notify { auto_run: true, .. }
|
||||
| Method::NotifyTracked { auto_run: true, .. }
|
||||
| Method::Resume
|
||||
| Method::Compact => WorkerExecutionRunState::Busy,
|
||||
Method::Shutdown => WorkerExecutionRunState::Stopped,
|
||||
@@ -1959,12 +1898,19 @@ where
|
||||
&& busy
|
||||
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_ok();
|
||||
let notification_request_id = input
|
||||
.submission_request_id
|
||||
.unwrap_or_else(protocol::new_submission_request_id);
|
||||
let result = self.send_method(
|
||||
WorkerExecutionOperation::Input,
|
||||
worker,
|
||||
Method::Notify {
|
||||
Method::NotifyTracked {
|
||||
notification_request_id: notification_request_id.clone(),
|
||||
message: input.content,
|
||||
auto_run: true,
|
||||
source: protocol::AuthenticatedInputSource::Backend {
|
||||
operation_id: notification_request_id,
|
||||
},
|
||||
},
|
||||
accepted_run_state,
|
||||
);
|
||||
@@ -1975,21 +1921,23 @@ where
|
||||
return result;
|
||||
}
|
||||
|
||||
if worker.shared_state.get_status() != WorkerStatus::Idle
|
||||
|| busy
|
||||
let is_user_submit = input.kind == WorkerInputKind::User;
|
||||
let status = worker.shared_state.get_status();
|
||||
let claimed_here = status == WorkerStatus::Idle
|
||||
&& busy
|
||||
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_err()
|
||||
{
|
||||
.is_ok();
|
||||
if !is_user_submit && !claimed_here {
|
||||
return WorkerExecutionResult::busy(
|
||||
WorkerExecutionOperation::Input,
|
||||
"Worker is already running; runtime adapter v0 does not queue input",
|
||||
"Worker is already running",
|
||||
);
|
||||
}
|
||||
|
||||
let (method, submission_id) = match input.kind {
|
||||
let (method, submission_request_id) = match input.kind {
|
||||
WorkerInputKind::User => {
|
||||
let Some(submission_id) = input
|
||||
.submission_id
|
||||
.submission_request_id
|
||||
.filter(|submission_id| !submission_id.trim().is_empty())
|
||||
else {
|
||||
busy.store(false, Ordering::SeqCst);
|
||||
@@ -1999,11 +1947,14 @@ where
|
||||
);
|
||||
};
|
||||
(
|
||||
Method::RunTracked {
|
||||
Method::SubmitTracked {
|
||||
submission_request_id: submission_id.clone(),
|
||||
input: input.segments.unwrap_or_else(|| {
|
||||
vec![Segment::text(input.content.trim().to_string())]
|
||||
}),
|
||||
submission_id: submission_id.clone(),
|
||||
source: protocol::AuthenticatedInputSource::Backend {
|
||||
operation_id: submission_id.clone(),
|
||||
},
|
||||
},
|
||||
Some(submission_id),
|
||||
)
|
||||
@@ -2021,21 +1972,22 @@ where
|
||||
),
|
||||
};
|
||||
let accepted_run_state = match method {
|
||||
Method::Run { .. }
|
||||
| Method::RunTracked { .. }
|
||||
Method::Submit { .. }
|
||||
| Method::SubmitTracked { .. }
|
||||
| Method::Notify { .. }
|
||||
| Method::NotifyTracked { .. }
|
||||
| Method::Compact => WorkerExecutionRunState::Busy,
|
||||
_ => WorkerExecutionRunState::Idle,
|
||||
};
|
||||
let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle;
|
||||
let waits_for_user_input_commit = submission_id.is_some();
|
||||
let waits_for_submission_acceptance = submission_request_id.is_some();
|
||||
|
||||
let result = if waits_for_user_input_commit {
|
||||
self.send_user_input_and_wait_for_commit(
|
||||
let result = if waits_for_submission_acceptance {
|
||||
self.send_submit_and_wait_for_acceptance(
|
||||
WorkerExecutionOperation::Input,
|
||||
worker,
|
||||
method,
|
||||
submission_id.expect("tracked Run has submission id"),
|
||||
submission_request_id.expect("Submit must have a submission request id"),
|
||||
accepted_run_state,
|
||||
)
|
||||
} else {
|
||||
@@ -2046,7 +1998,9 @@ where
|
||||
accepted_run_state,
|
||||
)
|
||||
};
|
||||
if accepted_is_idle || result.outcome != crate::execution::WorkerExecutionOutcome::Accepted
|
||||
if accepted_is_idle
|
||||
|| (claimed_here
|
||||
&& result.outcome != crate::execution::WorkerExecutionOutcome::Accepted)
|
||||
{
|
||||
busy.store(false, Ordering::SeqCst);
|
||||
}
|
||||
@@ -2116,8 +2070,12 @@ where
|
||||
}
|
||||
};
|
||||
|
||||
if let Method::Notify { auto_run, .. } = &method {
|
||||
let auto_run = *auto_run;
|
||||
if let Some(auto_run) = match &method {
|
||||
Method::Notify { auto_run, .. } | Method::NotifyTracked { auto_run, .. } => {
|
||||
Some(*auto_run)
|
||||
}
|
||||
_ => None,
|
||||
} {
|
||||
let status = worker.shared_state.get_status();
|
||||
let accepted_run_state = accepted_notify_run_state(status, auto_run);
|
||||
let claimed_here = status == WorkerStatus::Idle
|
||||
@@ -3400,6 +3358,60 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn running_worker_accepts_a_second_submit_as_queued() {
|
||||
let client = MockClient::sequential(vec![MockResponse::Hang(vec![])]);
|
||||
let runtime_base = tempfile::tempdir().unwrap();
|
||||
let cwd = tempfile::tempdir().unwrap();
|
||||
let store = tempfile::tempdir().unwrap();
|
||||
let factory = MockFactory {
|
||||
client,
|
||||
runtime_base: runtime_base.path().to_path_buf(),
|
||||
cwd: cwd.path().to_path_buf(),
|
||||
store_dir: store.path().join("sessions"),
|
||||
worker_metadata_dir: store.path().join("workers"),
|
||||
observed_cwds: Arc::new(Mutex::new(Vec::new())),
|
||||
observed_workspace_clients: Arc::new(Mutex::new(Vec::new())),
|
||||
};
|
||||
let backend = Arc::new(WorkerRuntimeExecutionBackend::new(factory).unwrap());
|
||||
let runtime =
|
||||
EmbeddedRuntime::with_execution_backend(RuntimeOptions::default(), backend).unwrap();
|
||||
runtime.store_config_bundle(test_bundle()).unwrap();
|
||||
let detail = runtime
|
||||
.create_worker(create_request("queued-submit"))
|
||||
.unwrap();
|
||||
|
||||
let mut first_input = WorkerInput::user("first");
|
||||
first_input.submission_request_id = Some("request-first".into());
|
||||
let first = runtime
|
||||
.send_input(&detail.worker_ref, first_input.clone())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
first.submission.as_ref().map(|ack| ack.disposition),
|
||||
Some(protocol::SubmissionDisposition::Started)
|
||||
);
|
||||
let retry = runtime.send_input(&detail.worker_ref, first_input).unwrap();
|
||||
assert_eq!(retry.submission, first.submission);
|
||||
let mut conflicting_retry = WorkerInput::user("different");
|
||||
conflicting_retry.submission_request_id = Some("request-first".into());
|
||||
assert!(
|
||||
runtime
|
||||
.send_input(&detail.worker_ref, conflicting_retry)
|
||||
.is_err(),
|
||||
"same request id with a different payload must fail"
|
||||
);
|
||||
|
||||
let mut second_input = WorkerInput::user("second");
|
||||
second_input.submission_request_id = Some("request-second".into());
|
||||
let second = runtime
|
||||
.send_input(&detail.worker_ref, second_input)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
second.submission.as_ref().map(|ack| ack.disposition),
|
||||
Some(protocol::SubmissionDisposition::Queued)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_with_initial_input_returns_after_session_commit() {
|
||||
let client = MockClient::new(simple_text_events());
|
||||
@@ -3449,8 +3461,10 @@ mod tests {
|
||||
};
|
||||
extensions
|
||||
.iter()
|
||||
.find(|extension| extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN)
|
||||
.and_then(|extension| extension.payload["submission_id"].as_str())
|
||||
.find(|extension| extension.domain == "worker.pending_activations.v1")
|
||||
.and_then(|extension| {
|
||||
extension.payload["receipts"][0]["submission_id"].as_str()
|
||||
})
|
||||
})
|
||||
.expect("committed input submission id");
|
||||
uuid::Uuid::parse_str(submission_id).expect("opaque submission id is a UUID");
|
||||
|
||||
@@ -20,6 +20,7 @@ protocol = { workspace = true, features = ["json-schema"] }
|
||||
client = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
reqwest = { version = "0.13", default-features = false, features = ["blocking", "native-tls"] }
|
||||
thiserror = { workspace = true }
|
||||
tokio = { workspace = true, features = ["fs", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "sync", "time"] }
|
||||
|
||||
@@ -101,7 +101,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
// Send a run method
|
||||
handle
|
||||
.send(Method::run_text("What is the capital of France?"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"What is the capital of France?",
|
||||
))
|
||||
.await?;
|
||||
|
||||
// Wait for completion
|
||||
|
||||
+793
-156
File diff suppressed because it is too large
Load Diff
@@ -1012,7 +1012,19 @@ async fn send_peer_notify(socket_path: &Path, message: String) -> io::Result<()>
|
||||
}
|
||||
|
||||
async fn send_notify(socket_path: &Path, message: String, auto_run: bool) -> io::Result<()> {
|
||||
connect_and_send(socket_path, &Method::Notify { message, auto_run }).await
|
||||
let notification_request_id = protocol::new_submission_request_id();
|
||||
connect_and_send(
|
||||
socket_path,
|
||||
&Method::NotifyTracked {
|
||||
notification_request_id: notification_request_id.clone(),
|
||||
message,
|
||||
auto_run,
|
||||
source: protocol::AuthenticatedInputSource::Backend {
|
||||
operation_id: notification_request_id,
|
||||
},
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn json_content<T: Serialize>(value: &T) -> Result<String, ToolError> {
|
||||
@@ -1482,6 +1494,7 @@ mod tests {
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -1517,6 +1530,7 @@ mod tests {
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -1536,7 +1550,10 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
let method = reader.next::<Method>().await.unwrap().unwrap();
|
||||
if let Method::Notify { message, auto_run } = method {
|
||||
if let Method::NotifyTracked {
|
||||
message, auto_run, ..
|
||||
} = method
|
||||
{
|
||||
assert!(auto_run);
|
||||
tx.send(message).await.unwrap();
|
||||
} else {
|
||||
@@ -1608,6 +1625,7 @@ mod tests {
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -1634,6 +1652,7 @@ mod tests {
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -1653,7 +1672,10 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
let method = reader.next::<Method>().await.unwrap().unwrap();
|
||||
if let Method::Notify { message, auto_run } = method {
|
||||
if let Method::NotifyTracked {
|
||||
message, auto_run, ..
|
||||
} = method
|
||||
{
|
||||
assert!(!auto_run);
|
||||
tx.send(message).await.unwrap();
|
||||
} else {
|
||||
@@ -1738,6 +1760,7 @@ mod tests {
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
@@ -1790,6 +1813,8 @@ mod tests {
|
||||
let _ = writer
|
||||
.write(&Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(
|
||||
),
|
||||
entries: Vec::new(),
|
||||
},
|
||||
greeting: protocol::Greeting {
|
||||
|
||||
@@ -803,7 +803,10 @@ mod tests {
|
||||
.collect();
|
||||
Ok(WorkerSessionCapture {
|
||||
segment_id: "segment".to_string(),
|
||||
session: protocol::SessionSnapshot { entries },
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -176,12 +176,23 @@ impl WorkerInterceptor {
|
||||
/// `Item::system_message`s reach the worker via
|
||||
/// `ContinueWith` / `pending_history_appends`, so on-disk order
|
||||
/// matches worker-history order.
|
||||
fn commit_system_items(&self, items: &[SystemItem]) -> Result<(), session_store::StoreError> {
|
||||
fn commit_system_items_with_extensions(
|
||||
&self,
|
||||
items: &[(
|
||||
SystemItem,
|
||||
Vec<session_store::SessionExtension>,
|
||||
Option<session_store::LoggedSessionHistoryOrigin>,
|
||||
)],
|
||||
) -> Result<(), session_store::StoreError> {
|
||||
let Some(writer) = self.log_writer.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
for item in items {
|
||||
let entry = writer.commit_system_item(item.clone())?;
|
||||
for (item, extensions, history_provenance) in items {
|
||||
let entry = writer.commit_system_item_with_extensions(
|
||||
item.clone(),
|
||||
extensions.clone(),
|
||||
history_provenance.clone(),
|
||||
)?;
|
||||
self.pending_committed_history
|
||||
.lock()
|
||||
.expect("pending committed history poisoned")
|
||||
@@ -190,6 +201,16 @@ impl WorkerInterceptor {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn commit_system_items(&self, items: &[SystemItem]) -> Result<(), session_store::StoreError> {
|
||||
self.commit_system_items_with_extensions(
|
||||
&items
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|item| (item, Vec::new(), None))
|
||||
.collect::<Vec<_>>(),
|
||||
)
|
||||
}
|
||||
|
||||
fn current_turn_index(&self) -> usize {
|
||||
self.next_turn_index
|
||||
.load(Ordering::Relaxed)
|
||||
@@ -327,7 +348,11 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
|
||||
projection_digest: projection.catalog_digest.clone(),
|
||||
logical_name: "internal.notify_wrapper".to_string(),
|
||||
};
|
||||
let mut system_items: Vec<SystemItem> = Vec::with_capacity(drained.len());
|
||||
let mut system_items: Vec<(
|
||||
SystemItem,
|
||||
Vec<session_store::SessionExtension>,
|
||||
Option<session_store::LoggedSessionHistoryOrigin>,
|
||||
)> = Vec::with_capacity(drained.len());
|
||||
let mut items: Vec<Item> = Vec::with_capacity(drained.len());
|
||||
for entry in &drained {
|
||||
let system_item = match build_system_item_with_provenance(
|
||||
@@ -345,9 +370,9 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
|
||||
}
|
||||
};
|
||||
items.push(system_item.to_history_item());
|
||||
system_items.push(system_item);
|
||||
system_items.push((system_item, entry.extensions(), entry.history_provenance()));
|
||||
}
|
||||
if let Err(error) = self.commit_system_items(&system_items) {
|
||||
if let Err(error) = self.commit_system_items_with_extensions(&system_items) {
|
||||
self.pending_notifies.requeue_front(drained);
|
||||
return Err(InterceptorError::new(
|
||||
InterceptorErrorCategory::Dependency,
|
||||
|
||||
@@ -25,7 +25,7 @@ use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use protocol::WorkerEvent;
|
||||
use session_store::SystemItem;
|
||||
use session_store::{LoggedSessionHistoryOrigin, SessionExtension, SystemItem};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::prompt::catalog::{CatalogError, PromptCatalog};
|
||||
@@ -41,8 +41,33 @@ const CAPACITY: usize = 128;
|
||||
/// is available.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PendingNotify {
|
||||
Notify { message: String, auto_run: bool },
|
||||
WorkerEvent { event: WorkerEvent },
|
||||
Notify {
|
||||
message: String,
|
||||
auto_run: bool,
|
||||
extensions: Vec<SessionExtension>,
|
||||
history_provenance: Option<LoggedSessionHistoryOrigin>,
|
||||
},
|
||||
WorkerEvent {
|
||||
event: WorkerEvent,
|
||||
},
|
||||
}
|
||||
|
||||
impl PendingNotify {
|
||||
pub(crate) fn extensions(&self) -> Vec<SessionExtension> {
|
||||
match self {
|
||||
PendingNotify::Notify { extensions, .. } => extensions.clone(),
|
||||
PendingNotify::WorkerEvent { .. } => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn history_provenance(&self) -> Option<LoggedSessionHistoryOrigin> {
|
||||
match self {
|
||||
PendingNotify::Notify {
|
||||
history_provenance, ..
|
||||
} => history_provenance.clone(),
|
||||
PendingNotify::WorkerEvent { .. } => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared, mutex-guarded buffer of pending entries.
|
||||
@@ -62,7 +87,46 @@ impl NotifyBuffer {
|
||||
/// oldest entry is dropped and a `tracing::warn` is emitted — the
|
||||
/// caller should never hit this in normal operation.
|
||||
pub fn push_notify(&self, message: String, auto_run: bool) {
|
||||
self.push_entry(PendingNotify::Notify { message, auto_run });
|
||||
self.push_entry(PendingNotify::Notify {
|
||||
message,
|
||||
auto_run,
|
||||
extensions: Vec::new(),
|
||||
history_provenance: None,
|
||||
});
|
||||
}
|
||||
|
||||
pub fn push_durable_notify(
|
||||
&self,
|
||||
message: String,
|
||||
auto_run: bool,
|
||||
history_provenance: LoggedSessionHistoryOrigin,
|
||||
extension: SessionExtension,
|
||||
) {
|
||||
self.push_entry(PendingNotify::Notify {
|
||||
message,
|
||||
auto_run,
|
||||
extensions: vec![extension],
|
||||
history_provenance: Some(history_provenance),
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn replace_durable_notification_extension(
|
||||
&self,
|
||||
extension: SessionExtension,
|
||||
) -> bool {
|
||||
let mut queue = self.inner.lock().expect("notify buffer poisoned");
|
||||
let Some(extensions) = queue.iter_mut().rev().find_map(|pending| match pending {
|
||||
PendingNotify::Notify {
|
||||
auto_run: false,
|
||||
extensions,
|
||||
..
|
||||
} if !extensions.is_empty() => Some(extensions),
|
||||
_ => None,
|
||||
}) else {
|
||||
return false;
|
||||
};
|
||||
*extensions = vec![extension];
|
||||
true
|
||||
}
|
||||
|
||||
/// Push a typed worker-event entry onto the queue.
|
||||
@@ -202,6 +266,8 @@ mod tests {
|
||||
let entry = PendingNotify::Notify {
|
||||
message: "hello".into(),
|
||||
auto_run: false,
|
||||
extensions: Vec::new(),
|
||||
history_provenance: None,
|
||||
};
|
||||
let catalog = PromptCatalog::builtins_only().unwrap();
|
||||
let item = build_system_item(&entry, &catalog).unwrap();
|
||||
|
||||
@@ -57,9 +57,9 @@ pub use session_history::{
|
||||
};
|
||||
pub use shared_state::WorkerSharedState;
|
||||
pub use worker::{
|
||||
LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
|
||||
WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient,
|
||||
WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution,
|
||||
WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, WorkspaceWorkerDiscoveryRequest,
|
||||
apply_worker_manifest, marker_workspace_client, unavailable_workspace_client,
|
||||
LocalWorkingDirectory, Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult,
|
||||
WorkerWorkspaceContext, WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError,
|
||||
WorkspacePromptCatalogResolution, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse,
|
||||
WorkspaceWorkerDiscoveryRequest, apply_worker_manifest, marker_workspace_client,
|
||||
unavailable_workspace_client,
|
||||
};
|
||||
|
||||
@@ -291,6 +291,7 @@ mod tests {
|
||||
prompt_provenance: None,
|
||||
},
|
||||
),
|
||||
extensions: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -72,6 +72,7 @@ mod tests {
|
||||
fn snapshot(entries: Vec<serde_json::Value>) -> Event {
|
||||
Event::Snapshot {
|
||||
session: protocol::SessionSnapshot {
|
||||
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
|
||||
entries: entries
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
|
||||
@@ -58,7 +58,7 @@ struct SubWorkerSpawnInput {
|
||||
/// a host path and grants no authority. When omitted, the Workdir root is used.
|
||||
#[serde(default)]
|
||||
cwd: Option<String>,
|
||||
/// First message sent to the spawned SubWorker via `Method::Run`.
|
||||
/// First message sent to the spawned SubWorker via `Method::Submit`.
|
||||
task: String,
|
||||
/// Allow rules delegated to the spawned SubWorker. Must be a subset of the
|
||||
/// spawner's explicit delegation authority; direct tool scope alone is not
|
||||
@@ -219,33 +219,50 @@ fn parse_spawn_profile_selector(raw: Option<&str>) -> Result<SpawnProfileSelecto
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum ParentNotificationTarget {
|
||||
Controller(mpsc::WeakSender<Method>),
|
||||
Buffer(crate::ipc::notify_buffer::NotifyBuffer),
|
||||
Controller {
|
||||
sender: mpsc::WeakSender<Method>,
|
||||
fallback: Arc<dyn Fn(Method) + Send + Sync>,
|
||||
},
|
||||
Durable(Arc<dyn Fn(Method) + Send + Sync>),
|
||||
}
|
||||
|
||||
impl ParentNotificationTarget {
|
||||
fn notify(&self, message: String, auto_run: bool) {
|
||||
pub(crate) fn with_controller_fallback(
|
||||
sender: mpsc::WeakSender<Method>,
|
||||
fallback: ParentNotificationTarget,
|
||||
) -> Self {
|
||||
let ParentNotificationTarget::Durable(fallback) = fallback else {
|
||||
unreachable!("controller fallback must use durable pending authority");
|
||||
};
|
||||
Self::Controller { sender, fallback }
|
||||
}
|
||||
|
||||
pub(crate) fn notify(&self, child_session_id: String, message: String, auto_run: bool) {
|
||||
let method = Method::NotifyTracked {
|
||||
notification_request_id: protocol::new_submission_request_id(),
|
||||
message,
|
||||
auto_run,
|
||||
source: protocol::AuthenticatedInputSource::SubWorker {
|
||||
session_id: child_session_id,
|
||||
},
|
||||
};
|
||||
match self {
|
||||
Self::Controller(parent_method_tx) => {
|
||||
let Some(parent_method_tx) = parent_method_tx.upgrade() else {
|
||||
tracing::warn!(
|
||||
"parent Worker controller closed before Internal SubWorker completion notification"
|
||||
);
|
||||
Self::Controller { sender, fallback } => {
|
||||
let Some(parent_method_tx) = sender.upgrade() else {
|
||||
fallback(method);
|
||||
return;
|
||||
};
|
||||
let fallback = fallback.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = parent_method_tx
|
||||
.send(Method::Notify { message, auto_run })
|
||||
.await
|
||||
{
|
||||
if let Err(error) = parent_method_tx.send(method).await {
|
||||
tracing::warn!(
|
||||
%error,
|
||||
"failed to notify parent Worker about Internal SubWorker completion"
|
||||
"failed to notify parent Controller; using durable pending authority"
|
||||
);
|
||||
fallback(error.0);
|
||||
}
|
||||
});
|
||||
}
|
||||
Self::Buffer(parent_notifies) => parent_notifies.push_notify(message, auto_run),
|
||||
Self::Durable(notify) => notify(method),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -550,7 +567,7 @@ impl Tool for SubWorkerSpawnTool {
|
||||
let message = format!(
|
||||
"SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions."
|
||||
);
|
||||
parent_notifications.notify(message, true);
|
||||
parent_notifications.notify(child_name.clone(), message, true);
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
@@ -1134,12 +1151,41 @@ enabled = false
|
||||
#[tokio::test]
|
||||
async fn parent_controller_notification_target_does_not_keep_channel_open() {
|
||||
let (parent_method_tx, mut parent_method_rx) = mpsc::channel(1);
|
||||
let target = ParentNotificationTarget::Controller(parent_method_tx.downgrade());
|
||||
let captured = Arc::new(std::sync::Mutex::new(false));
|
||||
let captured_for_fallback = captured.clone();
|
||||
let target = ParentNotificationTarget::with_controller_fallback(
|
||||
parent_method_tx.downgrade(),
|
||||
ParentNotificationTarget::Durable(Arc::new(move |_| {
|
||||
*captured_for_fallback.lock().unwrap() = true;
|
||||
})),
|
||||
);
|
||||
|
||||
drop(parent_method_tx);
|
||||
|
||||
assert!(parent_method_rx.recv().await.is_none());
|
||||
target.notify("late completion".to_string(), true);
|
||||
target.notify("child-session".into(), "late completion".to_string(), true);
|
||||
assert!(*captured.lock().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn durable_parent_notification_target_preserves_child_source() {
|
||||
let captured = Arc::new(std::sync::Mutex::new(None));
|
||||
let captured_for_target = captured.clone();
|
||||
let target = ParentNotificationTarget::Durable(Arc::new(move |method| {
|
||||
*captured_for_target.lock().unwrap() = Some(method);
|
||||
}));
|
||||
|
||||
target.notify("child-session".into(), "completed".into(), true);
|
||||
|
||||
assert!(matches!(
|
||||
captured.lock().unwrap().take(),
|
||||
Some(Method::NotifyTracked {
|
||||
message,
|
||||
auto_run: true,
|
||||
source: protocol::AuthenticatedInputSource::SubWorker { session_id },
|
||||
..
|
||||
}) if session_id == "child-session" && message == "completed"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1185,7 +1231,10 @@ enabled = false
|
||||
let tool = SubWorkerSpawnTool::new(
|
||||
"parent".into(),
|
||||
workspace_context,
|
||||
ParentNotificationTarget::Controller(parent_method_tx.downgrade()),
|
||||
ParentNotificationTarget::with_controller_fallback(
|
||||
parent_method_tx.downgrade(),
|
||||
ParentNotificationTarget::Durable(Arc::new(|_| {})),
|
||||
),
|
||||
runtime.path().to_path_buf(),
|
||||
bash_output_dir.clone(),
|
||||
workspace_root.clone(),
|
||||
@@ -1282,10 +1331,13 @@ enabled = false
|
||||
.expect("parent method channel remains open");
|
||||
assert!(matches!(
|
||||
completion,
|
||||
Method::Notify {
|
||||
Method::NotifyTracked {
|
||||
message,
|
||||
auto_run: true,
|
||||
} if message.contains("SubWorker `reviewer-child` turn ended with status Idle")
|
||||
source: protocol::AuthenticatedInputSource::SubWorker { session_id },
|
||||
..
|
||||
} if session_id == "reviewer-child"
|
||||
&& message.contains("SubWorker `reviewer-child` turn ended with status Idle")
|
||||
));
|
||||
assert!(!runtime.path().join("reviewer-child/sock").exists());
|
||||
|
||||
@@ -1434,7 +1486,10 @@ enabled = false
|
||||
let tool = SubWorkerSpawnTool::new(
|
||||
"parent".into(),
|
||||
workspace_context,
|
||||
ParentNotificationTarget::Controller(parent_method_tx.downgrade()),
|
||||
ParentNotificationTarget::with_controller_fallback(
|
||||
parent_method_tx.downgrade(),
|
||||
ParentNotificationTarget::Durable(Arc::new(|_| {})),
|
||||
),
|
||||
runtime.path().to_path_buf(),
|
||||
bash_output_dir.clone(),
|
||||
workspace_root.clone(),
|
||||
|
||||
+1670
-20
File diff suppressed because it is too large
Load Diff
@@ -630,7 +630,10 @@ async fn controller_compact_method_emits_start_and_done() {
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle
|
||||
.send(Method::run_text("seed history"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"seed history",
|
||||
))
|
||||
.await
|
||||
.expect("send run");
|
||||
loop {
|
||||
|
||||
@@ -617,7 +617,13 @@ async fn feature_flags_default_to_core_tool_surface_only() {
|
||||
let worker = make_worker(client).await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
|
||||
let request = wait_for_captured_request(&client_for_assert).await;
|
||||
@@ -672,7 +678,13 @@ permission = "write"
|
||||
let worker = make_worker_with_pwd_and_manifest(client, manifest).await.0;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
|
||||
let request = wait_for_captured_request(&client_for_assert).await;
|
||||
@@ -758,7 +770,13 @@ permission = "write"
|
||||
let worker = make_worker_with_pwd_and_manifest(client, &manifest).await.0;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
|
||||
let request = wait_for_captured_request(&client_for_assert).await;
|
||||
@@ -826,7 +844,13 @@ async fn builtin_orchestrator_exposes_worker_remove_and_workdir_delete() {
|
||||
.await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
let request = wait_for_captured_request(&client_for_assert).await;
|
||||
let installed = request_tool_names(&request);
|
||||
@@ -875,7 +899,13 @@ permission = "write"
|
||||
.0;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
|
||||
let request = wait_for_captured_request(&client_for_assert).await;
|
||||
@@ -928,7 +958,13 @@ permission = "write"
|
||||
)
|
||||
.await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
let request = wait_for_captured_request(&client_for_assert).await;
|
||||
let names = request_tool_names(&request);
|
||||
@@ -975,7 +1011,13 @@ async fn run_end_returns_to_idle_without_busy_status() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut saw_run_end = false;
|
||||
let mut saw_idle_status = false;
|
||||
@@ -1017,7 +1059,13 @@ async fn provider_stream_error_records_run_errored() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("ping")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"ping",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
|
||||
@@ -1066,7 +1114,10 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
|
||||
let mut events = handle.subscribe();
|
||||
|
||||
handle
|
||||
.send(Method::run_text("hello in-flight"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"hello in-flight",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
||||
@@ -1131,7 +1182,13 @@ async fn attach_snapshot_includes_current_status() {
|
||||
let worker = make_worker(client).await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Running).await;
|
||||
|
||||
let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path())
|
||||
@@ -1169,7 +1226,13 @@ async fn run_updates_shared_state_to_idle_after_completion() {
|
||||
let worker = make_worker(client).await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Wait for the run to complete
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
@@ -1183,7 +1246,13 @@ async fn run_populates_history() {
|
||||
let worker = make_worker(client).await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
|
||||
@@ -1201,7 +1270,13 @@ async fn events_are_broadcast() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("Hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut saw_turn_start = false;
|
||||
let mut saw_text_delta = false;
|
||||
@@ -1236,10 +1311,8 @@ async fn events_are_broadcast() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn double_run_returns_error() {
|
||||
// Keep the first turn in-flight until the test drops the handle. A
|
||||
// finite stream can finish before the second Method reaches the
|
||||
// controller in the full test suite, making this assertion racy.
|
||||
async fn submit_while_running_is_durably_queued() {
|
||||
// Keep the first turn in-flight until the second Submit is accepted.
|
||||
let events = vec![
|
||||
LlmEvent::text_block_start(0),
|
||||
LlmEvent::text_delta(0, "slow..."),
|
||||
@@ -1249,35 +1322,67 @@ async fn double_run_returns_error() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
// Send first run and wait until the controller has entered Running.
|
||||
handle.send(Method::run_text("first")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text("request-first", "first"))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Running).await;
|
||||
handle
|
||||
.send(Method::submit_text("request-second", "second"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Now the second run must be rejected by drive_turn's live Method arm.
|
||||
handle.send(Method::run_text("second")).await.unwrap();
|
||||
|
||||
// Look for the error event
|
||||
let mut saw_already_running = false;
|
||||
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
|
||||
loop {
|
||||
tokio::select! {
|
||||
event = rx.recv() => {
|
||||
match event {
|
||||
Ok(Event::Error { code, .. }) => {
|
||||
if code == worker::ErrorCode::AlreadyRunning {
|
||||
saw_already_running = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
let mut accepted = None;
|
||||
let mut pending_snapshot = None;
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
match tokio::time::timeout(std::time::Duration::from_millis(100), rx.recv()).await {
|
||||
Ok(Ok(Event::SubmissionAccepted {
|
||||
submission_request_id,
|
||||
disposition,
|
||||
..
|
||||
})) if submission_request_id == "request-second" => accepted = Some(disposition),
|
||||
Ok(Ok(Event::PendingSubmissionsChanged { pending }))
|
||||
if pending.submissions.len() == 1 =>
|
||||
{
|
||||
pending_snapshot = Some(pending)
|
||||
}
|
||||
_ = tokio::time::sleep_until(deadline) => break,
|
||||
Ok(Ok(Event::Error { code, message })) if code == worker::ErrorCode::AlreadyRunning => {
|
||||
panic!("Submit was busy-rejected: {message}")
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
if accepted.is_some() && pending_snapshot.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
assert!(saw_already_running, "should see already_running error");
|
||||
assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued));
|
||||
let pending_snapshot = pending_snapshot.expect("pending snapshot");
|
||||
assert_eq!(pending_snapshot.submissions.len(), 1);
|
||||
handle.send(Method::Pause).await.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Paused).await;
|
||||
handle
|
||||
.send(Method::ContinuePending {
|
||||
expected_revision: pending_snapshot.revision,
|
||||
expected_head_id: pending_snapshot.head_id.expect("pending head"),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let rejection = tokio::time::timeout(std::time::Duration::from_secs(1), async {
|
||||
loop {
|
||||
if let Ok(Event::Error { code, message }) = rx.recv().await
|
||||
&& code == worker::ErrorCode::InvalidRequest
|
||||
&& message.contains("requires an idle Worker")
|
||||
{
|
||||
break message;
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("paused ContinuePending rejection");
|
||||
assert!(rejection.contains("Resume or Cancel"));
|
||||
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Paused);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1365,7 +1470,8 @@ async fn run_with_paste_segment_inlines_content_and_emits_typed_user_message() {
|
||||
protocol::Segment::text(" thanks"),
|
||||
];
|
||||
handle
|
||||
.send(Method::Run {
|
||||
.send(Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: segments.clone(),
|
||||
})
|
||||
.await
|
||||
@@ -1437,7 +1543,13 @@ async fn run_with_resolvable_file_ref_attaches_system_message_after_user() {
|
||||
path: "notes.md".into(),
|
||||
},
|
||||
];
|
||||
handle.send(Method::Run { input: segments }).await.unwrap();
|
||||
handle
|
||||
.send(Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: segments,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Wait for the turn to complete.
|
||||
let mut rx = handle.subscribe();
|
||||
@@ -1485,7 +1597,8 @@ async fn run_with_file_ref_uses_manifest_file_upload_limit() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
|
||||
handle
|
||||
.send(Method::Run {
|
||||
.send(Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: vec![protocol::Segment::FileRef {
|
||||
path: "long.txt".into(),
|
||||
}],
|
||||
@@ -1538,7 +1651,13 @@ async fn run_with_unresolved_segment_emits_alert_and_placeholder() {
|
||||
path: "src/lib.rs".into(),
|
||||
},
|
||||
];
|
||||
handle.send(Method::Run { input: segments }).await.unwrap();
|
||||
handle
|
||||
.send(Method::Submit {
|
||||
submission_request_id: protocol::new_submission_request_id(),
|
||||
input: segments,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
|
||||
let mut saw_alert_for_file_ref = false;
|
||||
@@ -1586,6 +1705,7 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
|
||||
|
||||
handle
|
||||
.send(Method::Notify {
|
||||
notification_request_id: protocol::new_submission_request_id(),
|
||||
message: "turn finished".into(),
|
||||
auto_run: true,
|
||||
})
|
||||
@@ -1626,6 +1746,19 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
|
||||
saw_notify_in_mirror,
|
||||
"Method::Notify should commit a SystemItem::Notification entry; mirror = {entries:?}"
|
||||
);
|
||||
let queue_checkpoint_is_atomic = entries.iter().any(|entry| match entry {
|
||||
LogEntry::AnnotatedSystemItem { extensions, .. } => extensions.iter().any(|extension| {
|
||||
extension.domain == "worker.pending_activations.v1"
|
||||
&& extension.payload["pending_notifications"]
|
||||
.as_array()
|
||||
.is_some_and(Vec::is_empty)
|
||||
}),
|
||||
_ => false,
|
||||
});
|
||||
assert!(
|
||||
queue_checkpoint_is_atomic,
|
||||
"notification history and queue claim must share one log entry"
|
||||
);
|
||||
|
||||
// Exactly one request was made; it must contain the formatted
|
||||
// notification as one of the items (committed to history by
|
||||
@@ -1671,14 +1804,18 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
|
||||
let client_for_assert = client.clone();
|
||||
let worker = make_worker(client).await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
let notification_request_id = protocol::new_submission_request_id();
|
||||
|
||||
handle
|
||||
.send(Method::Notify {
|
||||
message: "progress snapshot".into(),
|
||||
auto_run: false,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
for _ in 0..2 {
|
||||
handle
|
||||
.send(Method::Notify {
|
||||
notification_request_id: notification_request_id.clone(),
|
||||
message: "progress snapshot".into(),
|
||||
auto_run: false,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
|
||||
@@ -1687,7 +1824,13 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
|
||||
"weak Notify must not stage RunForNotification while idle"
|
||||
);
|
||||
|
||||
handle.send(Method::run_text("continue")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"continue",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
|
||||
loop {
|
||||
if !client_for_assert.captured_requests().is_empty() {
|
||||
@@ -1867,9 +2010,16 @@ async fn notify_while_running_does_not_emit_already_running_error() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("start")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"start",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
handle
|
||||
.send(Method::Notify {
|
||||
notification_request_id: protocol::new_submission_request_id(),
|
||||
message: "ping".into(),
|
||||
auto_run: true,
|
||||
})
|
||||
@@ -1902,6 +2052,66 @@ async fn notify_while_running_does_not_emit_already_running_error() {
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn weak_notify_while_running_is_deduped_and_survives_until_next_submit() {
|
||||
let client = MockClient::sequential(vec![
|
||||
MockResponse::Hang(Vec::new()),
|
||||
MockResponse::Complete(simple_text_events()),
|
||||
]);
|
||||
let client_for_assert = client.clone();
|
||||
let worker = make_worker(client).await;
|
||||
let handle = spawn_controller(worker).await;
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"first",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Running).await;
|
||||
|
||||
let notification_request_id = protocol::new_submission_request_id();
|
||||
for _ in 0..2 {
|
||||
handle
|
||||
.send(Method::Notify {
|
||||
notification_request_id: notification_request_id.clone(),
|
||||
message: "durable weak notice".into(),
|
||||
auto_run: false,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
handle.send(Method::Cancel).await.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
|
||||
let mut rx = handle.subscribe();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"second",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
||||
loop {
|
||||
if matches!(rx.recv().await, Ok(Event::TurnEnd { .. })) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("second submit completes");
|
||||
|
||||
let requests = client_for_assert.captured_requests();
|
||||
let notice_count = requests[1]
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|item| item.as_text())
|
||||
.filter(|text| text.contains("durable weak notice"))
|
||||
.count();
|
||||
assert_eq!(notice_count, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn status_json_reflects_worker_name() {
|
||||
let client = MockClient::new(simple_text_events());
|
||||
@@ -1936,7 +2146,13 @@ async fn socket_run_receives_events() {
|
||||
let mut writer = JsonLineWriter::new(writer);
|
||||
|
||||
// Send run method via socket
|
||||
writer.write(&Method::run_text("Hello")).await.unwrap();
|
||||
writer
|
||||
.write(&Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"Hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Collect events
|
||||
let mut saw_turn_start = false;
|
||||
@@ -2243,7 +2459,13 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("hello")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"hello",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Wait for the partial text_delta to confirm the first stream is
|
||||
// live before we pause.
|
||||
@@ -2332,7 +2554,7 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
|
||||
assert!(!has_tool_call, "no orphan tool_call in history");
|
||||
}
|
||||
|
||||
/// Paused with an orphan `tool_use` in history + a fresh `Method::Run`
|
||||
/// Paused with an orphan `tool_use` in history + a fresh `Method::Submit`
|
||||
/// must produce a wire-valid next LLM request: the orphan is closed
|
||||
/// with a synthetic `tool_result`, a system note is inserted, and the
|
||||
/// new user input is appended.
|
||||
@@ -2369,7 +2591,13 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("first")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"first",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Wait for ToolCallDone — the ToolCall is committed to history
|
||||
// right before the Engine enters tool execution and pends.
|
||||
@@ -2400,7 +2628,13 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
|
||||
// `last_run_interrupted` and runs its interrupt-prep step, which
|
||||
// closes the orphan + injects a system note before the fresh user
|
||||
// message.
|
||||
handle.send(Method::run_text("new request")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"new request",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
|
||||
e,
|
||||
@@ -2531,7 +2765,13 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("first")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"first",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
|
||||
e,
|
||||
@@ -2599,7 +2839,10 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
|
||||
);
|
||||
|
||||
handle
|
||||
.send(Method::run_text("fresh request"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"fresh request",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
@@ -2688,7 +2931,13 @@ async fn empty_turn_cancel_rolls_back_submit_entries_and_emits_signal() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("rollback me")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"rollback me",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Running).await;
|
||||
handle.send(Method::Cancel).await.unwrap();
|
||||
|
||||
@@ -2721,7 +2970,10 @@ async fn empty_turn_pause_rolls_back_and_snapshot_does_not_restore_input() {
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle
|
||||
.send(Method::run_text("pause rollback"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"pause rollback",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Running).await;
|
||||
@@ -2755,7 +3007,13 @@ async fn empty_turn_rollback_removes_only_the_most_recent_turn() {
|
||||
let handle = spawn_controller(worker).await;
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle.send(Method::run_text("first kept")).await.unwrap();
|
||||
handle
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"first kept",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
|
||||
e,
|
||||
@@ -2769,7 +3027,10 @@ async fn empty_turn_rollback_removes_only_the_most_recent_turn() {
|
||||
wait_for_status(&handle, WorkerStatus::Idle).await;
|
||||
|
||||
handle
|
||||
.send(Method::run_text("second rolled back"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"second rolled back",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
wait_for_status(&handle, WorkerStatus::Running).await;
|
||||
@@ -2816,7 +3077,10 @@ async fn pause_after_assistant_token_does_not_rollback() {
|
||||
let mut rx = handle.subscribe();
|
||||
|
||||
handle
|
||||
.send(Method::run_text("keep this turn"))
|
||||
.send(Method::submit_text(
|
||||
protocol::new_submission_request_id(),
|
||||
"keep this turn",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
|
||||
@@ -538,7 +538,7 @@ fn initial_worker_input(segments: &[Segment]) -> Option<EmbeddedWorkerInput> {
|
||||
Some(EmbeddedWorkerInput {
|
||||
kind: EmbeddedWorkerInputKind::User,
|
||||
content: Segment::flatten_to_text(segments),
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: Some(segments.to_vec()),
|
||||
})
|
||||
}
|
||||
@@ -2706,7 +2706,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
|
||||
WorkerInputKind::RegisterPeer => EmbeddedWorkerInputKind::RegisterPeer,
|
||||
},
|
||||
content: request.content,
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: request.segments,
|
||||
};
|
||||
match self.runtime.send_input(&worker_ref, input) {
|
||||
@@ -3934,7 +3934,7 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
|
||||
WorkerInputKind::RegisterPeer => EmbeddedWorkerInputKind::RegisterPeer,
|
||||
},
|
||||
content: request.content,
|
||||
submission_id: None,
|
||||
submission_request_id: None,
|
||||
segments: request.segments,
|
||||
};
|
||||
match self.post_json::<_, RuntimeHttpWorkerInputResponse>(
|
||||
@@ -5195,7 +5195,7 @@ mod tests {
|
||||
"missing test context",
|
||||
);
|
||||
};
|
||||
let submission_id = input.submission_id.clone();
|
||||
let submission_request_id = input.submission_request_id.clone();
|
||||
let content = input.content;
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(std::time::Duration::from_millis(10));
|
||||
@@ -5212,11 +5212,13 @@ mod tests {
|
||||
status: protocol::WorkerStatus::Idle,
|
||||
});
|
||||
});
|
||||
if let Some(submission_id) = submission_id {
|
||||
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
|
||||
if let Some(submission_request_id) = submission_request_id {
|
||||
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
|
||||
worker_runtime::execution::WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Busy,
|
||||
submission_id,
|
||||
submission_request_id,
|
||||
uuid::Uuid::now_v7().to_string(),
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
} else {
|
||||
worker_runtime::execution::WorkerExecutionResult::accepted(
|
||||
|
||||
@@ -32,11 +32,13 @@ impl WorkerExecutionBackend for TestExecutionBackend {
|
||||
_handle: &WorkerExecutionHandle,
|
||||
input: worker_runtime::interaction::WorkerInput,
|
||||
) -> WorkerExecutionResult {
|
||||
if let Some(submission_id) = input.submission_id {
|
||||
WorkerExecutionResult::accepted_input_committed(
|
||||
if let Some(submission_request_id) = input.submission_request_id {
|
||||
WorkerExecutionResult::accepted_submission(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Busy,
|
||||
submission_id,
|
||||
submission_request_id,
|
||||
uuid::Uuid::now_v7().to_string(),
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
} else {
|
||||
WorkerExecutionResult::accepted(
|
||||
|
||||
@@ -8574,13 +8574,15 @@ async fn scoped_list_runtimes(
|
||||
|
||||
async fn scoped_workspace_protocol_ws(
|
||||
State(api): State<WorkspaceApi>,
|
||||
Extension(actor): Extension<RequestActor>,
|
||||
AxumPath(workspace_id): AxumPath<String>,
|
||||
ws: axum::extract::ws::WebSocketUpgrade,
|
||||
) -> std::result::Result<Response, Response> {
|
||||
validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?;
|
||||
let input_source = authenticated_browser_input_source(&actor);
|
||||
Ok(ws
|
||||
.on_upgrade(move |socket| {
|
||||
crate::workspace_subscription::serve_workspace_subscription(api, socket)
|
||||
crate::workspace_subscription::serve_workspace_subscription(api, socket, input_source)
|
||||
})
|
||||
.into_response())
|
||||
}
|
||||
@@ -9286,7 +9288,7 @@ async fn scoped_capture_worker_observation_session(
|
||||
return Err(ApiError::from(Error::UnknownWorker { worker: target }));
|
||||
}
|
||||
|
||||
let mut connection = connect_workspace_worker_protocol(&api, &target).await?;
|
||||
let mut connection = connect_workspace_worker_protocol(&api, &target, None).await?;
|
||||
let event = tokio::time::timeout(std::time::Duration::from_secs(10), connection.events.recv())
|
||||
.await
|
||||
.map_err(|_| {
|
||||
@@ -11495,6 +11497,7 @@ async fn scoped_cancel_runtime_worker(
|
||||
async fn scoped_worker_protocol_ws(
|
||||
ws: WebSocketUpgrade,
|
||||
State(api): State<WorkspaceApi>,
|
||||
Extension(actor): Extension<RequestActor>,
|
||||
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
|
||||
) -> Response {
|
||||
if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) {
|
||||
@@ -11502,6 +11505,7 @@ async fn scoped_worker_protocol_ws(
|
||||
}
|
||||
worker_protocol_ws(
|
||||
State(api),
|
||||
Extension(actor),
|
||||
AxumPath((path.worker.runtime_id, path.worker.worker_id)),
|
||||
ws,
|
||||
)
|
||||
@@ -13896,8 +13900,45 @@ async fn cancel_runtime_worker(
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
fn authenticated_browser_input_source(actor: &RequestActor) -> protocol::AuthenticatedInputSource {
|
||||
protocol::AuthenticatedInputSource::Account {
|
||||
account_id: actor.account_id.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn authorize_browser_worker_method(
|
||||
method: protocol::Method,
|
||||
source: &protocol::AuthenticatedInputSource,
|
||||
) -> std::result::Result<protocol::Method, &'static str> {
|
||||
match method {
|
||||
protocol::Method::Submit {
|
||||
submission_request_id,
|
||||
input,
|
||||
} => Ok(protocol::Method::SubmitTracked {
|
||||
submission_request_id,
|
||||
input,
|
||||
source: source.clone(),
|
||||
}),
|
||||
protocol::Method::Notify {
|
||||
notification_request_id,
|
||||
message,
|
||||
auto_run,
|
||||
} => Ok(protocol::Method::NotifyTracked {
|
||||
notification_request_id,
|
||||
message,
|
||||
auto_run,
|
||||
source: source.clone(),
|
||||
}),
|
||||
protocol::Method::SubmitTracked { .. } | protocol::Method::NotifyTracked { .. } => {
|
||||
Err("authenticated Worker input source is server-owned")
|
||||
}
|
||||
other => Ok(other),
|
||||
}
|
||||
}
|
||||
|
||||
async fn worker_protocol_ws(
|
||||
State(api): State<WorkspaceApi>,
|
||||
Extension(actor): Extension<RequestActor>,
|
||||
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
|
||||
ws: WebSocketUpgrade,
|
||||
) -> impl IntoResponse {
|
||||
@@ -13921,7 +13962,8 @@ async fn worker_protocol_ws(
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket))
|
||||
let input_source = authenticated_browser_input_source(&actor);
|
||||
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket, input_source))
|
||||
}
|
||||
|
||||
pub(crate) struct WorkspaceWorkerProtocolConnection {
|
||||
@@ -13932,6 +13974,7 @@ pub(crate) struct WorkspaceWorkerProtocolConnection {
|
||||
pub(crate) async fn connect_workspace_worker_protocol(
|
||||
api: &WorkspaceApi,
|
||||
worker: &RuntimeWorkerRef,
|
||||
input_source: Option<&protocol::AuthenticatedInputSource>,
|
||||
) -> Result<WorkspaceWorkerProtocolConnection> {
|
||||
let source = match api.observation_proxy.source(worker) {
|
||||
Ok(source) => source,
|
||||
@@ -13948,15 +13991,39 @@ pub(crate) async fn connect_workspace_worker_protocol(
|
||||
}
|
||||
};
|
||||
match source {
|
||||
RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await,
|
||||
RuntimeObservationSource::RemoteWs(config) => {
|
||||
connect_remote_worker_protocol(config, input_source).await
|
||||
}
|
||||
RuntimeObservationSource::Embedded(source) => {
|
||||
connect_embedded_worker_protocol(source).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_authenticated_input_source_header(
|
||||
headers: &mut HeaderMap,
|
||||
input_source: Option<&protocol::AuthenticatedInputSource>,
|
||||
) -> Result<()> {
|
||||
let Some(input_source) = input_source else {
|
||||
return Ok(());
|
||||
};
|
||||
let protocol::AuthenticatedInputSource::Account { account_id } = input_source else {
|
||||
return Err(Error::Config(
|
||||
"remote Worker protocol transport supports only Account input source".into(),
|
||||
));
|
||||
};
|
||||
headers.insert(
|
||||
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
|
||||
account_id.parse().map_err(|error| {
|
||||
Error::Config(format!("invalid authenticated Account identity: {error}"))
|
||||
})?,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn connect_remote_worker_protocol(
|
||||
config: RuntimeObservationSourceConfig,
|
||||
input_source: Option<&protocol::AuthenticatedInputSource>,
|
||||
) -> Result<WorkspaceWorkerProtocolConnection> {
|
||||
let mut request = config
|
||||
.endpoint
|
||||
@@ -13971,6 +14038,7 @@ async fn connect_remote_worker_protocol(
|
||||
})?,
|
||||
);
|
||||
}
|
||||
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
|
||||
let (socket, _) =
|
||||
connect_async(request)
|
||||
.await
|
||||
@@ -14048,13 +14116,17 @@ async fn connect_embedded_worker_protocol(
|
||||
Ok(WorkspaceWorkerProtocolConnection { methods, events })
|
||||
}
|
||||
|
||||
async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) {
|
||||
async fn worker_protocol_ws_session(
|
||||
source: RuntimeObservationSource,
|
||||
socket: WebSocket,
|
||||
input_source: protocol::AuthenticatedInputSource,
|
||||
) {
|
||||
match source {
|
||||
RuntimeObservationSource::RemoteWs(config) => {
|
||||
remote_worker_protocol_ws_session(config, socket).await;
|
||||
remote_worker_protocol_ws_session(config, socket, input_source).await;
|
||||
}
|
||||
RuntimeObservationSource::Embedded(source) => {
|
||||
embedded_worker_protocol_ws_session(source, socket).await;
|
||||
embedded_worker_protocol_ws_session(source, socket, input_source).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -14062,6 +14134,7 @@ async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: We
|
||||
async fn remote_worker_protocol_ws_session(
|
||||
config: RuntimeObservationSourceConfig,
|
||||
socket: WebSocket,
|
||||
input_source: protocol::AuthenticatedInputSource,
|
||||
) {
|
||||
let mut request = match config.endpoint.clone().into_client_request() {
|
||||
Ok(request) => request,
|
||||
@@ -14089,6 +14162,16 @@ async fn remote_worker_protocol_ws_session(
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Err(error) =
|
||||
insert_authenticated_input_source_header(request.headers_mut(), Some(&input_source))
|
||||
{
|
||||
let mut socket = socket;
|
||||
let event = protocol_error_event(format!(
|
||||
"failed to build authenticated Account identity header: {error}"
|
||||
));
|
||||
let _ = send_protocol_event(&mut socket, &event).await;
|
||||
return;
|
||||
}
|
||||
|
||||
let (upstream, _) = match connect_async(request).await {
|
||||
Ok(connection) => connection,
|
||||
@@ -14110,14 +14193,33 @@ async fn remote_worker_protocol_ws_session(
|
||||
inbound = client_stream.next() => {
|
||||
match inbound {
|
||||
Some(Ok(WsMessage::Text(text))) => {
|
||||
if upstream_sink.send(TungsteniteMessage::Text(text.to_string().into())).await.is_err() {
|
||||
let method = match protocol::stream::decode_method(text.as_ref()) {
|
||||
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
|
||||
Ok(method) => method,
|
||||
Err(message) => {
|
||||
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(message)) {
|
||||
let _ = client_sink.send(WsMessage::Text(event.into())).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
},
|
||||
Err(error) => {
|
||||
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(error.to_string())) {
|
||||
let _ = client_sink.send(WsMessage::Text(event.into())).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
};
|
||||
let Ok(method) = protocol::stream::encode_method(&method) else { break };
|
||||
if upstream_sink.send(TungsteniteMessage::Text(method.into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Some(Ok(WsMessage::Binary(binary))) => {
|
||||
if upstream_sink.send(TungsteniteMessage::Binary(binary.to_vec().into())).await.is_err() {
|
||||
break;
|
||||
Some(Ok(WsMessage::Binary(_))) => {
|
||||
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event("binary Worker methods are not accepted")) {
|
||||
let _ = client_sink.send(WsMessage::Text(event.into())).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
Some(Ok(WsMessage::Close(_))) | None => {
|
||||
let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await;
|
||||
@@ -14173,6 +14275,7 @@ async fn remote_worker_protocol_ws_session(
|
||||
async fn embedded_worker_protocol_ws_session(
|
||||
source: crate::observation::EmbeddedRuntimeObservationSource,
|
||||
mut socket: WebSocket,
|
||||
input_source: protocol::AuthenticatedInputSource,
|
||||
) {
|
||||
let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(
|
||||
source.clone(),
|
||||
@@ -14192,24 +14295,32 @@ async fn embedded_worker_protocol_ws_session(
|
||||
inbound = socket.next() => {
|
||||
match inbound {
|
||||
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
|
||||
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
|
||||
Ok(events) => {
|
||||
for event in events {
|
||||
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
|
||||
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
|
||||
Ok(events) => {
|
||||
for event in events {
|
||||
if !send_protocol_event(&mut socket, &event).await {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
let event = protocol_error_event(error.to_string());
|
||||
if !send_protocol_event(&mut socket, &event).await {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
let event = protocol_error_event(error.to_string());
|
||||
if !send_protocol_event(&mut socket, &event).await {
|
||||
return;
|
||||
}
|
||||
},
|
||||
Err(message) => {
|
||||
let event = protocol_error_event(message);
|
||||
let _ = send_protocol_event(&mut socket, &event).await;
|
||||
return;
|
||||
}
|
||||
},
|
||||
Err(error) => {
|
||||
let event =
|
||||
protocol_error_event(format!("malformed protocol method frame: {error}"));
|
||||
let event = protocol_error_event(format!(
|
||||
"malformed protocol method frame: {error}"
|
||||
));
|
||||
if !send_protocol_event(&mut socket, &event).await {
|
||||
return;
|
||||
}
|
||||
@@ -16219,6 +16330,48 @@ mod tests {
|
||||
&tail[..end]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn browser_worker_methods_receive_server_owned_account_source() {
|
||||
let source = protocol::AuthenticatedInputSource::Account {
|
||||
account_id: "account-1".into(),
|
||||
};
|
||||
let method = authorize_browser_worker_method(
|
||||
protocol::Method::Submit {
|
||||
submission_request_id: "request-1".into(),
|
||||
input: vec![protocol::Segment::text("hello")],
|
||||
},
|
||||
&source,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
method,
|
||||
protocol::Method::SubmitTracked {
|
||||
source: protocol::AuthenticatedInputSource::Account { ref account_id },
|
||||
..
|
||||
} if account_id == "account-1"
|
||||
));
|
||||
assert!(authorize_browser_worker_method(method, &source).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remote_worker_protocol_header_preserves_authenticated_account_source() {
|
||||
let mut headers = HeaderMap::new();
|
||||
insert_authenticated_input_source_header(
|
||||
&mut headers,
|
||||
Some(&protocol::AuthenticatedInputSource::Account {
|
||||
account_id: "account-1".into(),
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER)
|
||||
.unwrap(),
|
||||
"account-1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
|
||||
let source = include_str!("server.rs");
|
||||
@@ -18468,7 +18621,7 @@ mod tests {
|
||||
.get(handle.worker_ref())
|
||||
.cloned()
|
||||
.expect("execution context");
|
||||
let submission_id = input.submission_id.clone();
|
||||
let submission_request_id = input.submission_request_id.clone();
|
||||
let content = input.content.clone();
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(std::time::Duration::from_millis(25));
|
||||
@@ -18476,11 +18629,13 @@ mod tests {
|
||||
text: format!("server companion echoed: {content}"),
|
||||
});
|
||||
});
|
||||
if let Some(submission_id) = submission_id {
|
||||
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
|
||||
if let Some(submission_request_id) = submission_request_id {
|
||||
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
|
||||
worker_runtime::execution::WorkerExecutionOperation::Input,
|
||||
worker_runtime::execution::WorkerExecutionRunState::Idle,
|
||||
submission_id,
|
||||
submission_request_id,
|
||||
uuid::Uuid::now_v7().to_string(),
|
||||
protocol::SubmissionDisposition::Started,
|
||||
)
|
||||
} else {
|
||||
worker_runtime::execution::WorkerExecutionResult::accepted(
|
||||
@@ -27069,6 +27224,16 @@ mod tests {
|
||||
(runtime, worker_ref, endpoint)
|
||||
}
|
||||
|
||||
fn test_browser_request_actor() -> RequestActor {
|
||||
RequestActor {
|
||||
user_id: "test-user".into(),
|
||||
account_id: format!("account-{TEST_WORKSPACE_ID}"),
|
||||
handle: "test".into(),
|
||||
display_name: "Test".into(),
|
||||
auth_method: ActorAuthMethod::BrowserSession,
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_workspace_proxy(
|
||||
source: RuntimeObservationSourceConfig,
|
||||
) -> (String, tempfile::TempDir) {
|
||||
@@ -27087,11 +27252,8 @@ mod tests {
|
||||
.unwrap();
|
||||
let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let app_addr = app_listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(app_listener, build_inner_router(api))
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
|
||||
tokio::spawn(async move { axum::serve(app_listener, app).await.unwrap() });
|
||||
(
|
||||
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"),
|
||||
dir,
|
||||
@@ -27103,7 +27265,8 @@ mod tests {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = build_inner_router(test_api(dir.path()).await);
|
||||
let app = build_inner_router(test_api(dir.path()).await)
|
||||
.layer(Extension(test_browser_request_actor()));
|
||||
let server = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app).await;
|
||||
});
|
||||
@@ -27151,7 +27314,7 @@ mod tests {
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = build_inner_router(api);
|
||||
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
|
||||
let server = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app).await;
|
||||
});
|
||||
|
||||
@@ -11,7 +11,9 @@ use tokio::sync::mpsc;
|
||||
use worker_runtime::identity::RuntimeWorkerRef;
|
||||
|
||||
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
|
||||
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
|
||||
use crate::server::{
|
||||
WorkspaceApi, authorize_browser_worker_method, connect_workspace_worker_protocol,
|
||||
};
|
||||
use crate::store::WorkspaceResourceKind;
|
||||
|
||||
const OUTBOUND_CAPACITY: usize = 256;
|
||||
@@ -21,7 +23,11 @@ struct ActiveSubscription {
|
||||
methods: Option<mpsc::Sender<protocol::Method>>,
|
||||
}
|
||||
|
||||
pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) {
|
||||
pub(crate) async fn serve_workspace_subscription(
|
||||
api: WorkspaceApi,
|
||||
socket: WebSocket,
|
||||
input_source: protocol::AuthenticatedInputSource,
|
||||
) {
|
||||
let broker = api.runtime_subscription_broker().clone();
|
||||
let (mut socket_sender, mut socket_receiver) = socket.split();
|
||||
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
||||
@@ -85,7 +91,13 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
|
||||
runtime_id: Some(runtime_id),
|
||||
} => {
|
||||
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
|
||||
match connect_workspace_worker_protocol(&api, &worker).await {
|
||||
match connect_workspace_worker_protocol(
|
||||
&api,
|
||||
&worker,
|
||||
Some(&input_source),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(connection) => {
|
||||
let methods = connection.methods.clone();
|
||||
let task = tokio::spawn(run_worker_protocol(
|
||||
@@ -153,7 +165,12 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
|
||||
else {
|
||||
break;
|
||||
};
|
||||
if methods.send(message.method).await.is_err() {
|
||||
let Ok(method) =
|
||||
authorize_browser_worker_method(message.method, &input_source)
|
||||
else {
|
||||
break;
|
||||
};
|
||||
if methods.send(method).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user