From 1e2fdb46e546b0b16c62b9546fad77eb45877675 Mon Sep 17 00:00:00 2001 From: Daniel Vigovszky Date: Fri, 4 Sep 2026 12:42:07 +0000 Subject: [PATCH] Reduce copies in agent invocation hot path Amp-Thread-ID: https://ampcode.com/threads/T-01a06bd5-2a53-73c4-90a0-52a8883f5f02 Co-authored-by: Amp --- golem-common/src/schema/agent/mod.rs | 38 +- golem-common/src/schema/agent/tests.rs | 97 ++- golem-worker-executor/src/durable_host/mod.rs | 2 +- .../src/grpc/invocation_session.rs | 719 ++++++++++++++++-- golem-worker-executor/src/grpc/mod.rs | 2 +- .../src/services/oplog/mod.rs | 46 +- .../src/services/oplog/tests.rs | 144 ++++ golem-worker-executor/src/services/rpc.rs | 16 +- golem-worker-executor/src/worker/mod.rs | 65 +- 9 files changed, 999 insertions(+), 130 deletions(-) diff --git a/golem-common/src/schema/agent/mod.rs b/golem-common/src/schema/agent/mod.rs index dfa322f944..026b9e0441 100644 --- a/golem-common/src/schema/agent/mod.rs +++ b/golem-common/src/schema/agent/mod.rs @@ -40,7 +40,7 @@ use crate::schema::metadata::MetadataEnvelope; use crate::schema::schema_type::{NamedFieldType, SchemaType}; use crate::schema::schema_value::SchemaValue; use crate::schema::validation::placement::validate_agent_type_placement; -use crate::schema::validation::value::validate_value; +use crate::schema::validation::value::{validate_record_fields, validate_value}; use golem_schema_derive::{FromSchema, IntoSchema}; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -74,6 +74,12 @@ pub fn json_input_schema_value_to_typed_schema_value( graph: &SchemaGraph, input_schema: &InputSchema, ) -> Result { + let result_graph = projected_input_graph(graph, input_schema); + validate_input_value(&result_graph, &value)?; + Ok(TypedSchemaValue::new(result_graph, value)) +} + +fn projected_input_graph(graph: &SchemaGraph, input_schema: &InputSchema) -> SchemaGraph { // Only user-supplied fields are part of the caller's input; auto-injected // fields (e.g. the principal) are filled by the host out of band and are // not present in the incoming value, so they are excluded from the record @@ -89,18 +95,36 @@ pub fn json_input_schema_value_to_typed_schema_value( }) .collect(); let root = SchemaType::record(fields); - let result_graph = SchemaGraph { + SchemaGraph { defs: reachable_defs(graph, &root), root, + } +} + +fn validate_input_value(graph: &SchemaGraph, value: &SchemaValue) -> Result<(), String> { + let validation = match (&graph.root, value) { + ( + SchemaType::Record { + fields: field_types, + .. + }, + SchemaValue::Record { fields: values }, + ) => validate_record_fields( + graph, + field_types + .iter() + .map(|field| (field.name.as_str(), &field.body)), + values, + ), + _ => validate_value(graph, &graph.root, value), }; - validate_value(&result_graph, &result_graph.root, &value).map_err(|errors| { + validation.map_err(|errors| { errors .into_iter() .map(|err| err.to_string()) .collect::>() .join("; ") - })?; - Ok(TypedSchemaValue::new(result_graph, value)) + }) } pub use crate::schema::graph::reachable_defs; @@ -379,8 +403,8 @@ impl AgentMethodSchema { /// Validates the caller-supplied parameter record against this method's /// input schema and the owning agent's graph. pub fn validate_input(&self, graph: &SchemaGraph, input: &SchemaValue) -> Result<(), String> { - json_input_schema_value_to_typed_schema_value(input.clone(), graph, &self.input_schema) - .map(|_| ()) + let input_graph = projected_input_graph(graph, &self.input_schema); + validate_input_value(&input_graph, input) } /// Returns whether a caller-supplied input or the output of this method can diff --git a/golem-common/src/schema/agent/tests.rs b/golem-common/src/schema/agent/tests.rs index 1111a5e96c..d252103d4a 100644 --- a/golem-common/src/schema/agent/tests.rs +++ b/golem-common/src/schema/agent/tests.rs @@ -18,13 +18,14 @@ use crate::schema::agent::{ AgentConfigDeclarationSchema, AgentConstructorSchema, AgentDependencySchema, AgentMethodSchema, AgentTypeSchema, AutoInjectedKind, FieldSource, InputSchema, NamedField, OutputSchema, ParsedAgentId, contains_stream_in_graph, json_input_schema_value_to_typed_schema_value, - typed_schema_value_with_projected_defs, + reachable_defs, typed_schema_value_with_projected_defs, }; use crate::schema::graph::{SchemaGraph, SchemaTypeDef, TypedSchemaValue}; use crate::schema::metadata::{MetadataEnvelope, TypeId}; use crate::schema::schema_type::{NamedFieldType, SchemaType, SecretSpec, VariantCaseType}; use crate::schema::schema_value::{SchemaValue, SecretValuePayload, VariantValuePayload}; use crate::schema::stream::SchemaValueStream; +use crate::schema::validation::validate_value; use proptest::prelude::*; use serde_json::json; use test_r::test; @@ -775,15 +776,103 @@ fn method_input_validation_accepts_a_live_stream_handle() { )], OutputSchema::Unit, ); + let stream = SchemaValueStream::from_host_endpoint(()); let input = SchemaValue::Record { - fields: vec![SchemaValue::Stream(SchemaValueStream::from_host_endpoint( - (), - ))], + fields: vec![SchemaValue::Stream(stream.clone())], }; method .validate_input(&SchemaGraph::empty(), &input) .unwrap(); + assert!(stream.is_present()); +} + +#[test] +fn method_input_validation_matches_owned_conversion() { + fn assert_matches_owned_conversion( + method: &AgentMethodSchema, + graph: &SchemaGraph, + input: SchemaValue, + ) { + let fields = method + .input_schema + .fields() + .iter() + .filter(|field| matches!(field.source, FieldSource::UserSupplied)) + .map(|field| NamedFieldType { + name: field.name.clone(), + body: field.schema.clone(), + metadata: field.metadata.clone(), + }) + .collect(); + let root = SchemaType::record(fields); + let projected_graph = SchemaGraph { + defs: reachable_defs(graph, &root), + root, + }; + let expected = + validate_value(&projected_graph, &projected_graph.root, &input).map_err(|errors| { + errors + .into_iter() + .map(|error| error.to_string()) + .collect::>() + .join("; ") + }); + assert_eq!(method.validate_input(graph, &input), expected); + } + + let scalar_method = method( + "compare", + vec![ + NamedField::user_supplied("name", SchemaType::string()), + NamedField::auto_injected( + "principal", + AutoInjectedKind::Principal, + SchemaType::string(), + ), + NamedField::user_supplied("count", SchemaType::u64()), + ], + OutputSchema::Unit, + ); + for input in [ + SchemaValue::Record { + fields: vec![SchemaValue::String("ok".to_string()), SchemaValue::U64(1)], + }, + SchemaValue::Record { + fields: vec![SchemaValue::String("missing-count".to_string())], + }, + SchemaValue::Record { + fields: vec![ + SchemaValue::U64(1), + SchemaValue::String("wrong".to_string()), + ], + }, + SchemaValue::String("not-a-record".to_string()), + ] { + assert_matches_owned_conversion(&scalar_method, &SchemaGraph::empty(), input); + } + + let graph = registry(vec![proj_def( + "Payload", + SchemaType::record(vec![proj_field("value", SchemaType::string())]), + )]); + let method = method( + "compare-ref", + vec![NamedField::user_supplied( + "payload", + SchemaType::ref_to(TypeId::new("Payload")), + )], + OutputSchema::Unit, + ); + assert_matches_owned_conversion( + &method, + &graph, + SchemaValue::Record { + fields: vec![SchemaValue::Record { + fields: vec![SchemaValue::String("ok".to_string())], + }], + }, + ); } #[test] diff --git a/golem-worker-executor/src/durable_host/mod.rs b/golem-worker-executor/src/durable_host/mod.rs index f54d73bd1c..ee6ba77758 100644 --- a/golem-worker-executor/src/durable_host/mod.rs +++ b/golem-worker-executor/src/durable_host/mod.rs @@ -4489,7 +4489,7 @@ impl InvocationHooks for DurableWorkerCtx { _ => {} } - let (start_index, _) = self + let start_index = self .public_state .worker() .oplog() diff --git a/golem-worker-executor/src/grpc/invocation_session.rs b/golem-worker-executor/src/grpc/invocation_session.rs index 3329cb31ce..deba6566ca 100644 --- a/golem-worker-executor/src/grpc/invocation_session.rs +++ b/golem-worker-executor/src/grpc/invocation_session.rs @@ -19,7 +19,7 @@ use crate::durable_host::durable_session::{ use crate::durable_host::durable_stream::ProducerRegistrationRequestV1; use crate::durable_host::stream_session::{ decode_recursive_stream_value, decode_recursive_stream_value_with_schema, - remap_recursive_stream_references, + encode_recursive_stream_value_with_schema, }; use crate::grpc::invocation::{CanStartWorker, from_proto_invocation_context}; use crate::services::{HasAll, HasComponentService, HasSchedulerService, UsesAllDeps}; @@ -126,6 +126,18 @@ struct AcceptedInvocation { durable_replayed: bool, } +struct TransportStreamId(u64); + +pub(crate) fn decode_invocation_input( + input: golem_api_grpc::proto::golem::schema::SchemaValue, +) -> Result { + decode_recursive_stream_value(input, |stream_id, _| { + Ok(SchemaValueStream::from_host_endpoint(TransportStreamId( + stream_id, + ))) + }) +} + async fn detach_durable_attachment(streams: Option) { if let Some(streams) = streams && let Err(error) = streams.detach_current().await @@ -273,7 +285,7 @@ impl + UsesAllDeps + Send + Sync + &self, request: &InvocationStart, method_parameters: Option, - durable_input: Option, + input_encoded_len: Option, acceptance_committed: tokio::sync::oneshot::Sender<()>, accepted: tokio::sync::oneshot::Sender, ) -> Result { @@ -400,15 +412,6 @@ impl + UsesAllDeps + Send + Sync + .limit_invocation_context_stack_depth(from_proto_invocation_context(&request.context)); let worker_creation_principal = principal.clone(); - let invocation = AgentInvocation::AgentMethod { - idempotency_key: ik.clone(), - method_name: method_name.clone(), - input: method_parameters.clone(), - invocation_context, - principal, - scope_card, - }; - match mode { golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await => { let worker = self @@ -466,10 +469,18 @@ impl + UsesAllDeps + Send + Sync + }; let accepted_revision = (worker.agent_mode() == AgentMode::Ephemeral).then_some(component.revision); + let invocation = AgentInvocation::AgentMethod { + idempotency_key: ik.clone(), + method_name: method_name.clone(), + input: method_parameters, + invocation_context, + principal, + scope_card, + }; let mut invocation_output = if streaming { - let durable_input = durable_input.ok_or_else(|| { + let input_encoded_len = input_encoded_len.ok_or_else(|| { WorkerExecutorError::invalid_request( - "durable streaming invocation is missing its canonical input", + "durable streaming invocation is missing its encoded input length", ) })?; let request = build_durable_streaming_request( @@ -477,8 +488,8 @@ impl + UsesAllDeps + Send + Sync + &component.metadata, component.revision, worker.get_initial_worker_metadata().fingerprint, - invocation.clone(), - durable_input, + invocation, + input_encoded_len, acceptance_committed, self.services .config() @@ -556,6 +567,14 @@ impl + UsesAllDeps + Send + Sync + "live streams require an attached Await invocation session", )); } + let invocation = AgentInvocation::AgentMethod { + idempotency_key: ik.clone(), + method_name, + input: method_parameters, + invocation_context, + principal, + scope_card, + }; match schedule_at { Some(scheduled_time) => { @@ -700,7 +719,7 @@ impl + UsesAllDeps + Send + Sync + return; } let first = first.request.expect("validated request has a payload"); - let start = match first { + let mut start = match first { invocation_request::Request::Start(start) => start, invocation_request::Request::ResumeAttach(resume) => { self.run_resumed_agent_session(resume, inbound, outward, state) @@ -729,15 +748,13 @@ impl + UsesAllDeps + Send + Sync + } }); + let input_encoded_len = start.input.as_ref().map(Message::encoded_len); let input = if start.mode() == golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup { Ok(None) } else { - match start.input.clone() { - Some(input) => decode_recursive_stream_value(input, |_, _| { - Ok(SchemaValueStream::from_host_endpoint(())) - }) - .map(Some), + match start.input.take() { + Some(input) => decode_invocation_input(input).map(Some), None => Err("invocation start has no input".to_string()), } }; @@ -763,7 +780,7 @@ impl + UsesAllDeps + Send + Sync + let invocation = self.invoke_agent_internal( &start, input, - start.input.clone(), + input_encoded_len, acceptance_committed_tx, accepted_tx, ); @@ -1791,7 +1808,7 @@ pub(crate) fn build_durable_streaming_request( component_revision: ComponentRevision, callee_fingerprint: golem_common::model::AgentFingerprint, invocation: AgentInvocation, - durable_input: golem_api_grpc::proto::golem::schema::SchemaValue, + input_encoded_len: usize, acceptance_committed: tokio::sync::oneshot::Sender<()>, live_join_buffer_events: usize, ) -> Result { @@ -1858,6 +1875,14 @@ pub(crate) fn build_durable_streaming_request( }) .collect(), ); + let invocation_input = match &invocation { + AgentInvocation::AgentMethod { input, .. } => input, + _ => { + return Err(WorkerExecutorError::invalid_request( + "durable streaming request requires an agent method invocation", + )); + } + }; let session_key = StreamInvocationIdV1 { callee_environment_id: environment_id, callee: callee.clone(), @@ -1871,7 +1896,7 @@ pub(crate) fn build_durable_streaming_request( attachment_id, role: SessionStreamRoleV1::Input, }; - if durable_input.encoded_len() > MAX_DURABLE_STREAM_ITEM_SIZE { + if input_encoded_len > MAX_DURABLE_STREAM_ITEM_SIZE { return Err(WorkerExecutorError::invalid_request( "ResourceExhausted: durable invocation input exceeds the 16 MiB logical value limit", )); @@ -1903,51 +1928,59 @@ pub(crate) fn build_durable_streaming_request( } let mut input_element_types = Vec::new(); let mut canonical_foreign_mappings = Vec::new(); - decode_recursive_stream_value_with_schema( - durable_input.clone(), + let mut canonical_handle_index = 0u64; + let canonical_input = encode_recursive_stream_value_with_schema( + invocation_input, &agent_type.schema, &input_root, - |transport_stream_id, path| { - let element = stream_element_schema(&agent_type.schema, &input_root, path)?; - let element_schema_fingerprint = schema_fingerprint_v1(&agent_type.schema, element) - .map_err(|error| error.to_string())?; - input_element_types.push(( - transport_stream_id, - element.cloned().unwrap_or_else(SchemaType::u8), - )); - if foreign_mappings.is_empty() { - registrations.push(( + |stream, path| { + let transport_stream_id = stream + .with_host_endpoint::(|stream_id| stream_id.0)?; + let element = stream_element_schema(&agent_type.schema, &input_root, path)?; + let element_schema_fingerprint = schema_fingerprint_v1(&agent_type.schema, element) + .map_err(|error| error.to_string())?; + input_element_types.push(( transport_stream_id, - ProducerRegistrationRequestV1 { - coordinate: StreamRegistrationCoordinateV1::Root { - invocation_id: session_key.clone(), - root_kind: StreamRootKindV1::MethodInput, - recursive_value_path: path.to_vec(), - }, - source_kind: StreamSourceKindV1::ExternalInlineInput, - source_invocation: session_key.clone(), - component_revision, - element_schema_fingerprint, - session_mapping: Some(session_mapping.clone()), - }, + element.cloned().unwrap_or_else(SchemaType::u8), )); - } else { - let mapping = foreign_by_transport - .get(&transport_stream_id) - .ok_or_else(|| { - format!( - "durable invocation input references unmapped transport stream {transport_stream_id}" - ) - })? - .clone(); - if mapping.handle.element_schema_fingerprint != element_schema_fingerprint { - return Err(format!( - "durable invocation input stream {transport_stream_id} has the wrong schema fingerprint" + if foreign_mappings.is_empty() { + registrations.push(( + transport_stream_id, + ProducerRegistrationRequestV1 { + coordinate: StreamRegistrationCoordinateV1::Root { + invocation_id: session_key.clone(), + root_kind: StreamRootKindV1::MethodInput, + recursive_value_path: path.to_vec(), + }, + source_kind: StreamSourceKindV1::ExternalInlineInput, + source_invocation: session_key.clone(), + component_revision, + element_schema_fingerprint, + session_mapping: Some(session_mapping.clone()), + }, )); + } else { + let mapping = foreign_by_transport + .get(&transport_stream_id) + .ok_or_else(|| { + format!( + "durable invocation input references unmapped transport stream {transport_stream_id}" + ) + })? + .clone(); + if mapping.handle.element_schema_fingerprint != element_schema_fingerprint { + return Err(format!( + "durable invocation input stream {transport_stream_id} has the wrong schema fingerprint" + )); + } + canonical_foreign_mappings.push(mapping); } - canonical_foreign_mappings.push(mapping); - } - Ok(SchemaValueStream::from_host_endpoint(())) + + let index = canonical_handle_index; + canonical_handle_index = canonical_handle_index + .checked_add(1) + .ok_or_else(|| "durable input handle index overflow".to_string())?; + Ok(index) }, ) .map_err(WorkerExecutorError::invalid_request)?; @@ -1976,15 +2009,6 @@ pub(crate) fn build_durable_streaming_request( "ResourceExhausted: durable invocation input materializes more than 256 streams", )); } - let mut canonical_handle_index = 0u64; - let canonical_input = remap_recursive_stream_references(durable_input, |_, _| { - let index = canonical_handle_index; - canonical_handle_index = canonical_handle_index - .checked_add(1) - .ok_or_else(|| "durable input handle index overflow".to_string())?; - Ok(index) - }) - .map_err(WorkerExecutorError::invalid_request)?; if canonical_handle_index != expected_handle_count as u64 { return Err(WorkerExecutorError::invalid_request( "durable invocation input stream topology changed during canonicalization", @@ -2749,9 +2773,15 @@ async fn route_durable_request( #[cfg(test)] mod freshness_tests { use super::{ - AcceptanceRace, decode_invocation_freshness_disposition, effective_session_identity, - is_attachment_termination, pre_acceptance_rejection_reason, race_invocation_acceptance, - require_expected_callee_fingerprint, + AcceptanceRace, TransportStreamId, build_durable_streaming_request, + decode_invocation_freshness_disposition, decode_invocation_input, + effective_session_identity, is_attachment_termination, pre_acceptance_rejection_reason, + race_invocation_acceptance, require_expected_callee_fingerprint, + }; + use crate::durable_host::durable_session::durable_stream_mapping_to_proto; + use crate::durable_host::stream_session::{ + decode_recursive_stream_value, decode_recursive_stream_value_with_schema, + encode_recursive_stream_value_with_schema, remap_recursive_stream_references, }; use futures::future; use golem_api_grpc::proto::golem::account::PlanId; @@ -2759,13 +2789,546 @@ mod freshness_tests { AuthCtx, AuthEffectiveSurface, UserAuthCtx, auth_ctx, }; use golem_api_grpc::proto::golem::common::{AccountId, Uuid}; - use golem_api_grpc::proto::golem::worker::InvocationRejectionReason; - use golem_common::model::AgentFingerprint; - use golem_common::model::agent::InvocationFreshnessDisposition; + use golem_api_grpc::proto::golem::schema::{ + BinaryValue, SchemaValue as ProtoSchemaValue, schema_value as proto_schema_value, + }; + use golem_api_grpc::proto::golem::worker::{InvocationRejectionReason, InvocationStart}; + use golem_common::base_model::Empty; + use golem_common::base_model::agent::Snapshotting; + use golem_common::base_model::component::{ComponentId, ComponentRevision}; + use golem_common::base_model::component_metadata::{ComponentMetadata, KnownExports}; + use golem_common::base_model::durable_stream::{ + DURABLE_STREAM_FORMAT_VERSION, DurableStreamHandleV1, MAX_DURABLE_STREAM_ITEM_SIZE, + SessionStreamRoleV1, StreamId, StreamInvocationIdV1, StreamMapSideV1, + StreamSessionMappingRecordV1, StreamValuePathStepV1, + }; + use golem_common::base_model::environment::EnvironmentId; + use golem_common::model::agent::{ + AgentMode, AgentTypeName, InvocationFreshnessDisposition, Principal, + }; + use golem_common::model::invocation_context::InvocationContextStack; + use golem_common::model::{AgentFingerprint, AgentId, AgentInvocation, IdempotencyKey}; + use golem_common::schema::SchemaValue; + use golem_common::schema::agent::{ + AgentConstructorSchema, AgentMethodSchema, AgentTypeSchema, InputSchema, NamedField, + OutputSchema, + }; + use golem_schema::schema::schema_value::{ResultValuePayload, UnionValuePayload}; + use golem_schema::schema::{ + DiscriminatorRule, FieldDiscriminator, NamedFieldType, ResultSpec, SchemaFingerprintV1, + SchemaGraph, SchemaType, SchemaTypeDef, SchemaValueStream, TypeId, UnionBranch, UnionSpec, + VariantCaseType, schema_fingerprint_v1, + }; use golem_service_base::error::worker_executor::WorkerExecutorError; + use prost::Message; + use std::collections::BTreeMap; use std::task::Poll; use test_r::test; + #[test] + fn invocation_input_decode_moves_binary_payload() { + let bytes = vec![7; 1024]; + let bytes_ptr = bytes.as_ptr(); + let input = ProtoSchemaValue { + value: Some(proto_schema_value::Value::BinaryValue(BinaryValue { + bytes, + mime_type: None, + })), + }; + + let decoded = decode_invocation_input(input).unwrap(); + let SchemaValue::Binary(binary) = decoded else { + panic!("expected binary input") + }; + + assert_eq!(binary.bytes.as_ptr(), bytes_ptr); + } + + #[test] + fn one_pass_input_canonicalization_matches_two_pass_traversal() { + fn field(name: &str, body: SchemaType) -> NamedFieldType { + NamedFieldType { + name: name.to_string(), + body, + metadata: Default::default(), + } + } + + let map_type_id = TypeId::new("stream-map"); + let union = SchemaType::union(UnionSpec { + branches: vec![ + UnionBranch { + tag: "plain".to_string(), + body: SchemaType::record(vec![field("kind", SchemaType::string())]), + discriminator: DiscriminatorRule::FieldEquals(FieldDiscriminator { + field_name: "kind".to_string(), + literal: Some("plain".to_string()), + }), + metadata: Default::default(), + }, + UnionBranch { + tag: "stream".to_string(), + body: SchemaType::record(vec![ + field("kind", SchemaType::string()), + field("values", SchemaType::stream(Some(SchemaType::u64()))), + ]), + discriminator: DiscriminatorRule::FieldEquals(FieldDiscriminator { + field_name: "kind".to_string(), + literal: Some("stream".to_string()), + }), + metadata: Default::default(), + }, + ], + }); + let root = SchemaType::record(vec![ + field("mapping", SchemaType::ref_to(map_type_id.clone())), + field( + "variant", + SchemaType::variant(vec![ + VariantCaseType { + name: "none".to_string(), + payload: None, + metadata: Default::default(), + }, + VariantCaseType { + name: "some".to_string(), + payload: Some(SchemaType::option(SchemaType::stream(Some( + SchemaType::string(), + )))), + metadata: Default::default(), + }, + ]), + ), + field( + "result", + SchemaType::result(ResultSpec { + ok: Some(Box::new(SchemaType::u64())), + err: Some(Box::new(SchemaType::stream(Some(SchemaType::string())))), + }), + ), + field("union", union), + ]); + let graph = SchemaGraph { + defs: vec![SchemaTypeDef { + id: map_type_id, + name: None, + body: SchemaType::map( + SchemaType::stream(Some(SchemaType::u64())), + SchemaType::stream(Some(SchemaType::string())), + ), + }], + root: root.clone(), + }; + let input = SchemaValue::Record { + fields: vec![ + SchemaValue::Map { + entries: vec![( + SchemaValue::Stream(SchemaValueStream::from_host_endpoint(101_u64)), + SchemaValue::Stream(SchemaValueStream::from_host_endpoint(7_u64)), + )], + }, + SchemaValue::Variant(golem_schema::schema::schema_value::VariantValuePayload { + case: 1, + payload: Some(Box::new(SchemaValue::Option { + inner: Some(Box::new(SchemaValue::Stream( + SchemaValueStream::from_host_endpoint(55_u64), + ))), + })), + }), + SchemaValue::Result(ResultValuePayload::Err { + value: Some(Box::new(SchemaValue::Stream( + SchemaValueStream::from_host_endpoint(300_u64), + ))), + }), + SchemaValue::Union(UnionValuePayload { + tag: "stream".to_string(), + body: Box::new(SchemaValue::Record { + fields: vec![ + SchemaValue::String("stream".to_string()), + SchemaValue::Stream(SchemaValueStream::from_host_endpoint(2_u64)), + ], + }), + }), + ], + }; + let original = + encode_recursive_stream_value_with_schema(&input, &graph, &root, |stream, _| { + stream.with_host_endpoint::(|stream_id| *stream_id) + }) + .unwrap(); + + let mut old_observations = Vec::new(); + decode_recursive_stream_value_with_schema( + original.clone(), + &graph, + &root, + |stream_id, path| { + old_observations.push((stream_id, path.to_vec())); + Ok(SchemaValueStream::from_host_endpoint(())) + }, + ) + .unwrap(); + let mut old_index = 0_u64; + let old_canonical = remap_recursive_stream_references(original.clone(), |_, _| { + let index = old_index; + old_index += 1; + Ok(index) + }) + .unwrap(); + + let decoded = decode_invocation_input(original).unwrap(); + let canonicalize = |decoded: &SchemaValue| { + let mut observations = Vec::new(); + let mut index = 0_u64; + let canonical = encode_recursive_stream_value_with_schema( + decoded, + &graph, + &root, + |stream, path| { + let stream_id = stream + .with_host_endpoint::(|stream_id| stream_id.0)?; + observations.push((stream_id, path.to_vec())); + let canonical = index; + index += 1; + Ok(canonical) + }, + ) + .unwrap(); + (canonical, observations) + }; + let (new_canonical, new_observations) = canonicalize(&decoded); + + assert_eq!(new_canonical, old_canonical); + assert_eq!(new_observations, old_observations); + assert_eq!( + new_observations, + vec![ + ( + 101, + vec![ + StreamValuePathStepV1::RecordField(0), + StreamValuePathStepV1::MapEntry { + index: 0, + side: StreamMapSideV1::Key, + }, + ], + ), + ( + 7, + vec![ + StreamValuePathStepV1::RecordField(0), + StreamValuePathStepV1::MapEntry { + index: 0, + side: StreamMapSideV1::Value, + }, + ], + ), + ( + 55, + vec![ + StreamValuePathStepV1::RecordField(1), + StreamValuePathStepV1::VariantCasePayload(1), + StreamValuePathStepV1::OptionSome, + ], + ), + ( + 300, + vec![ + StreamValuePathStepV1::RecordField(2), + StreamValuePathStepV1::ResultErr, + ], + ), + ( + 2, + vec![ + StreamValuePathStepV1::RecordField(3), + StreamValuePathStepV1::UnionBranch(1), + StreamValuePathStepV1::RecordField(1), + ], + ), + ] + ); + assert_eq!(canonicalize(&decoded).0, old_canonical); + } + + fn builder_fixture() -> (InvocationStart, ComponentMetadata, AgentInvocation) { + let mapping_type = SchemaType::map( + SchemaType::stream(Some(SchemaType::u64())), + SchemaType::stream(Some(SchemaType::string())), + ); + let tail_type = SchemaType::stream(Some(SchemaType::bool())); + let method = AgentMethodSchema { + name: "run".to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![ + NamedField::user_supplied("mapping", mapping_type), + NamedField::user_supplied("tail", tail_type), + ]), + output_schema: OutputSchema::Unit, + http_endpoint: Vec::new(), + read_only: None, + }; + let metadata = ComponentMetadata::from_parts( + KnownExports::default(), + Vec::new(), + None, + None, + vec![AgentTypeSchema { + type_name: AgentTypeName("test-agent".to_string()), + description: String::new(), + source_language: String::new(), + schema: SchemaGraph::empty(), + constructor: AgentConstructorSchema { + name: None, + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(Vec::new()), + }, + methods: vec![method], + dependencies: Vec::new(), + mode: AgentMode::Durable, + http_mount: None, + snapshotting: Snapshotting::Disabled(Empty {}), + config: Vec::new(), + }], + BTreeMap::new(), + ); + let component_id = ComponentId(uuid::Uuid::from_u128(1)); + let environment_id = EnvironmentId(uuid::Uuid::from_u128(2)); + let callee = AgentId { + component_id, + agent_id: "test-agent()".to_string(), + }; + let idempotency_key = IdempotencyKey::new("test-key".to_string()); + let fingerprint = AgentFingerprint(uuid::Uuid::from_u128(3)); + let request = InvocationStart { + agent_id: Some(callee.into()), + method_name: Some("run".to_string()), + idempotency_key: Some(idempotency_key.clone().into()), + environment_id: Some(environment_id.into()), + attempt_id: Some(uuid::Uuid::new_v4().into()), + expected_callee_fingerprint: Some(fingerprint.0.into()), + ..Default::default() + }; + let invocation = AgentInvocation::AgentMethod { + idempotency_key, + method_name: "run".to_string(), + input: SchemaValue::Record { + fields: vec![ + SchemaValue::Map { + entries: vec![( + SchemaValue::Stream(SchemaValueStream::from_host_endpoint( + TransportStreamId(101), + )), + SchemaValue::Stream(SchemaValueStream::from_host_endpoint( + TransportStreamId(7), + )), + )], + }, + SchemaValue::Stream(SchemaValueStream::from_host_endpoint(TransportStreamId( + 55, + ))), + ], + }, + invocation_context: InvocationContextStack::fresh(), + principal: Principal::anonymous(), + scope_card: None, + }; + (request, metadata, invocation) + } + + fn foreign_mapping( + transport_stream_id: u64, + element_schema_fingerprint: SchemaFingerprintV1, + request: &InvocationStart, + ) -> StreamSessionMappingRecordV1 { + let environment_id: EnvironmentId = request.environment_id.unwrap().try_into().unwrap(); + let callee: AgentId = request.agent_id.clone().unwrap().try_into().unwrap(); + let callee_fingerprint = + AgentFingerprint(request.expected_callee_fingerprint.unwrap().into()); + let idempotency_key = request.idempotency_key.clone().unwrap().into(); + StreamSessionMappingRecordV1 { + transport_stream_id, + handle: DurableStreamHandleV1 { + format_version: DURABLE_STREAM_FORMAT_VERSION, + stream_id: StreamId(uuid::Uuid::from_u128( + 1_000 + u128::from(transport_stream_id), + )), + producer_environment_id: environment_id, + producer: callee.clone(), + expected_producer_fingerprint: callee_fingerprint, + source_invocation: StreamInvocationIdV1 { + callee_environment_id: environment_id, + callee, + callee_fingerprint, + idempotency_key, + }, + component_revision: ComponentRevision::INITIAL, + element_schema_fingerprint, + }, + role: SessionStreamRoleV1::Input, + } + } + + #[test] + fn durable_request_builder_preserves_size_precedence_and_stream_metadata_order() { + let (mut request, metadata, invocation) = builder_fixture(); + request.durable_input_mappings.push(Default::default()); + let (acceptance_committed, _) = tokio::sync::oneshot::channel(); + let error = match build_durable_streaming_request( + &request, + &metadata, + ComponentRevision::INITIAL, + AgentFingerprint(uuid::Uuid::from_u128(3)), + invocation, + MAX_DURABLE_STREAM_ITEM_SIZE + 1, + acceptance_committed, + 8, + ) { + Ok(_) => panic!("oversized input must be rejected"), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("input exceeds the 16 MiB logical value limit") + ); + + let (request, metadata, invocation) = builder_fixture(); + let (acceptance_committed, _) = tokio::sync::oneshot::channel(); + let built = build_durable_streaming_request( + &request, + &metadata, + ComponentRevision::INITIAL, + AgentFingerprint(uuid::Uuid::from_u128(3)), + invocation, + 1, + acceptance_committed, + 8, + ) + .unwrap(); + let expected_paths = [ + vec![ + StreamValuePathStepV1::RecordField(0), + StreamValuePathStepV1::MapEntry { + index: 0, + side: StreamMapSideV1::Key, + }, + ], + vec![ + StreamValuePathStepV1::RecordField(0), + StreamValuePathStepV1::MapEntry { + index: 0, + side: StreamMapSideV1::Value, + }, + ], + vec![StreamValuePathStepV1::RecordField(1)], + ]; + assert_eq!( + built + .registrations + .iter() + .map(|(transport_stream_id, registration)| { + let path = match ®istration.coordinate { + super::StreamRegistrationCoordinateV1::Root { + recursive_value_path, + .. + } => recursive_value_path.clone(), + _ => panic!("expected a root registration"), + }; + ( + *transport_stream_id, + path, + registration.element_schema_fingerprint, + ) + }) + .collect::>(), + vec![ + ( + 101, + expected_paths[0].clone(), + schema_fingerprint_v1(&SchemaGraph::empty(), Some(&SchemaType::u64())).unwrap(), + ), + ( + 7, + expected_paths[1].clone(), + schema_fingerprint_v1(&SchemaGraph::empty(), Some(&SchemaType::string())) + .unwrap(), + ), + ( + 55, + expected_paths[2].clone(), + schema_fingerprint_v1(&SchemaGraph::empty(), Some(&SchemaType::bool())) + .unwrap(), + ), + ] + ); + assert_eq!( + built.input_element_types, + vec![ + (101, SchemaType::u64()), + (7, SchemaType::string()), + (55, SchemaType::bool()), + ] + ); + let canonical = + ProtoSchemaValue::decode(built.attempt.invocation.invocation_value.as_slice()).unwrap(); + let mut canonical_ids = Vec::new(); + decode_recursive_stream_value(canonical, |stream_id, _| { + canonical_ids.push(stream_id); + Ok(SchemaValueStream::from_host_endpoint(())) + }) + .unwrap(); + assert_eq!(canonical_ids, vec![0, 1, 2]); + + let (mut request, metadata, invocation) = builder_fixture(); + let fingerprints = [ + schema_fingerprint_v1(&SchemaGraph::empty(), Some(&SchemaType::u64())).unwrap(), + schema_fingerprint_v1(&SchemaGraph::empty(), Some(&SchemaType::string())).unwrap(), + schema_fingerprint_v1(&SchemaGraph::empty(), Some(&SchemaType::bool())).unwrap(), + ]; + let mappings = [ + foreign_mapping(55, fingerprints[2], &request), + foreign_mapping(101, fingerprints[0], &request), + foreign_mapping(7, fingerprints[1], &request), + ]; + request.durable_input_mappings = mappings + .iter() + .map(|mapping| durable_stream_mapping_to_proto(mapping, None)) + .collect(); + let (acceptance_committed, _) = tokio::sync::oneshot::channel(); + let built = build_durable_streaming_request( + &request, + &metadata, + ComponentRevision::INITIAL, + AgentFingerprint(uuid::Uuid::from_u128(3)), + invocation, + 1, + acceptance_committed, + 8, + ) + .unwrap(); + + assert!(built.registrations.is_empty()); + assert_eq!( + built + .foreign_mappings + .iter() + .map(|mapping| mapping.transport_stream_id) + .collect::>(), + vec![101, 7, 55] + ); + assert_eq!( + built.input_element_types, + vec![ + (101, SchemaType::u64()), + (7, SchemaType::string()), + (55, SchemaType::bool()), + ] + ); + } + #[test] fn output_pump_distinguishes_attachment_termination_from_protocol_failure() { assert!(is_attachment_termination( diff --git a/golem-worker-executor/src/grpc/mod.rs b/golem-worker-executor/src/grpc/mod.rs index 3c85d16fc1..19140abe24 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -15,7 +15,7 @@ mod invocation; mod invocation_session; -pub(crate) use invocation_session::build_durable_streaming_request; +pub(crate) use invocation_session::{build_durable_streaming_request, decode_invocation_input}; use crate::durable_host::agent_monomorphization_context; use crate::grpc::invocation::{CanStartWorker, from_proto_invocation_context}; diff --git a/golem-worker-executor/src/services/oplog/mod.rs b/golem-worker-executor/src/services/oplog/mod.rs index a9d086edeb..cab4e66b6a 100644 --- a/golem-worker-executor/src/services/oplog/mod.rs +++ b/golem-worker-executor/src/services/oplog/mod.rs @@ -707,6 +707,17 @@ pub trait OplogOps: Oplog { Ok(payload) } + /// Uploads an owned oplog payload and moves it into the in-memory cache. + async fn upload_payload_owned( + &self, + data: T, + ) -> Result, String> { + let bytes = serialize(&data)?; + let raw_payload = self.upload_raw_payload(bytes).await?; + let payload = raw_payload.into_payload_with_cache(Arc::new(data))?; + Ok(payload) + } + /// Downloads a big oplog payload by its reference async fn download_payload( &self, @@ -851,32 +862,41 @@ pub trait OplogOps: Oplog { invocation: AgentInvocation, wallet_pin: InvocationWalletPin, ) -> Result { - self.add_agent_invocation_started_with_index(invocation, wallet_pin) - .await - .map(|(_, entry)| entry) + let entry = self + .agent_invocation_started_entry(invocation, wallet_pin) + .await?; + self.add(entry.clone()).await; + Ok(entry) } async fn add_agent_invocation_started_with_index( &self, invocation: AgentInvocation, wallet_pin: InvocationWalletPin, - ) -> Result<(OplogIndex, OplogEntry), String> { + ) -> Result { + let entry = self + .agent_invocation_started_entry(invocation, wallet_pin) + .await?; + Ok(self.add(entry).await) + } + + async fn agent_invocation_started_entry( + &self, + invocation: AgentInvocation, + wallet_pin: InvocationWalletPin, + ) -> Result { let (idempotency_key, invocation_payload, ctx) = invocation.into_parts(); - let payload = self.upload_payload(&invocation_payload).await?; - let trace_id = ctx.trace_id.clone(); - let trace_states = ctx.trace_states.clone(); + let payload = self.upload_payload_owned(invocation_payload).await?; let invocation_context = ctx.to_oplog_data(); - let entry = OplogEntry::AgentInvocationStarted { + Ok(OplogEntry::AgentInvocationStarted { timestamp: Timestamp::now_utc(), idempotency_key, payload, - trace_id, - trace_states, + trace_id: ctx.trace_id, + trace_states: ctx.trace_states, invocation_context, wallet_pin: Some(wallet_pin), - }; - let index = self.add(entry.clone()).await; - Ok((index, entry)) + }) } async fn add_agent_invocation_finished( diff --git a/golem-worker-executor/src/services/oplog/tests.rs b/golem-worker-executor/src/services/oplog/tests.rs index a29dfd5fb2..9d26f4912d 100644 --- a/golem-worker-executor/src/services/oplog/tests.rs +++ b/golem-worker-executor/src/services/oplog/tests.rs @@ -2787,6 +2787,59 @@ async fn completed_host_call_response_upload_failure_writes_no_start(_tracing: & )); } +#[test] +async fn owned_invocation_payload_upload_failure_writes_no_entry(_tracing: &Tracing) { + let indexed_storage = Arc::new(InMemoryIndexedStorage::new()); + let blob_storage = Arc::new(ReadCountingBlobStorage::failing_on_put(1)); + let oplog_service = PrimaryOplogService::new( + indexed_storage, + blob_storage, + 1, + 1, + 100, + RetryConfig::default(), + ) + .await; + let account_id = AccountId::new(); + let environment_id = EnvironmentId::new(); + let agent_id = AgentId { + component_id: ComponentId(Uuid::new_v4()), + agent_id: "owned-invocation-upload-failure".to_string(), + }; + let owned_agent_id = OwnedAgentId::new(environment_id, &agent_id); + let oplog = oplog_service + .open( + &owned_agent_id, + AgentMode::Durable, + None, + make_agent_metadata(agent_id, account_id, environment_id), + default_last_known_status(), + default_execution_status(AgentMode::Durable), + ) + .await; + let before = oplog.current_oplog_index().await; + + let result = oplog + .add_agent_invocation_started_with_index( + AgentInvocation::AgentMethod { + idempotency_key: IdempotencyKey::fresh(), + method_name: "large-input".to_string(), + input: SchemaValue::Binary(BinaryValuePayload { + bytes: vec![1_u8; 1024], + mime_type: None, + }), + invocation_context: InvocationContextStack::fresh_rounded(), + principal: Principal::anonymous(), + scope_card: None, + }, + invocation_wallet_pin(), + ) + .await; + + assert!(result.is_err()); + assert_eq!(oplog.current_oplog_index().await, before); +} + #[test] async fn entries_with_large_payload(_tracing: &Tracing) { let indexed_storage = Arc::new(InMemoryIndexedStorage::new()); @@ -5611,6 +5664,97 @@ async fn scan_for_component_with_no_workers_terminates_immediately(_tracing: &Tr } } +#[test] +async fn owned_payload_upload_moves_cache_and_roundtrips_without_it(_tracing: &Tracing) { + let indexed_storage = Arc::new(InMemoryIndexedStorage::new()); + let blob_storage = Arc::new(InMemoryBlobStorage::new()); + let oplog_service = PrimaryOplogService::new( + indexed_storage, + blob_storage, + 1, + 1, + 100, + RetryConfig::default(), + ) + .await; + let account_id = AccountId::new(); + let environment_id = EnvironmentId::new(); + let agent_id = AgentId { + component_id: ComponentId(Uuid::new_v4()), + agent_id: "owned-payload".to_string(), + }; + let owned_agent_id = OwnedAgentId::new(environment_id, &agent_id); + let oplog = oplog_service + .open( + &owned_agent_id, + AgentMode::Durable, + None, + make_agent_metadata(agent_id.clone(), account_id, environment_id), + default_last_known_status(), + default_execution_status(AgentMode::Durable), + ) + .await; + + let inline = vec![1_u8; 8]; + let inline_ptr = inline.as_ptr(); + let inline_without_cache = match oplog.upload_payload_owned(inline).await.unwrap() { + OplogPayload::SerializedInline { + bytes, + cached: Some(cached), + } => { + assert_eq!(cached.as_ptr(), inline_ptr); + OplogPayload::SerializedInline { + bytes, + cached: None, + } + } + other => panic!("expected an inline payload with a cache, got {other:?}"), + }; + + let external = vec![2_u8; 1024]; + let external_ptr = external.as_ptr(); + let external_without_cache = match oplog.upload_payload_owned(external).await.unwrap() { + OplogPayload::External { + payload_id, + md5_hash, + cached: Some(cached), + } => { + assert_eq!(cached.as_ptr(), external_ptr); + OplogPayload::External { + payload_id, + md5_hash, + cached: None, + } + } + other => panic!("expected an external payload with a cache, got {other:?}"), + }; + + let reopened = oplog_service + .open( + &owned_agent_id, + AgentMode::Durable, + None, + make_agent_metadata(agent_id, account_id, environment_id), + default_last_known_status(), + default_execution_status(AgentMode::Durable), + ) + .await; + assert_eq!( + reopened + .download_payload::>(inline_without_cache) + .await + .unwrap(), + vec![1_u8; 8] + ); + assert_eq!( + reopened + .download_payload::>(external_without_cache) + .await + .unwrap(), + vec![2_u8; 1024] + ); +} + /// A large request reserved with [`OplogOps::add_start_with_reserved_payload`] is stored externally, /// and its deferred blob upload is made durable by the leaf oplog's commit barrier even when the /// caller never awaits the returned [`PendingUpload`]. diff --git a/golem-worker-executor/src/services/rpc.rs b/golem-worker-executor/src/services/rpc.rs index d99445766b..82e5c23b00 100644 --- a/golem-worker-executor/src/services/rpc.rs +++ b/golem-worker-executor/src/services/rpc.rs @@ -17,9 +17,8 @@ use super::direct_invocation_auth::DirectInvocationAuthService; use super::environment_state::EnvironmentStateService; use super::file_loader::FileLoader; use super::{HasAgentWebhooksService, HasEnvironmentStateService, HasWebSocketConnectionPool}; -use crate::durable_host::stream_session::decode_recursive_stream_value; use crate::durable_host::websocket::WebSocketConnectionPool; -use crate::grpc::build_durable_streaming_request; +use crate::grpc::{build_durable_streaming_request, decode_invocation_input}; use crate::services::events::Events; use crate::services::oplog::plugin::OplogProcessorPlugin; use crate::services::resource_limits::ResourceLimits; @@ -65,9 +64,9 @@ use golem_common::model::{ AgentFingerprint, AgentId, AgentInvocation, AgentInvocationResult, IdempotencyKey, OwnedAgentId, }; use golem_common::schema::SchemaValue; -use golem_schema::schema::SchemaValueStream; use golem_service_base::error::worker_executor::WorkerExecutorError; use golem_service_base::model::auth::AuthCtx; +use prost::Message; use std::collections::HashMap; use std::fmt::{Display, Formatter}; use std::future::Future; @@ -1475,10 +1474,9 @@ impl Rpc for DirectWorkerInvocationRpc { Some(status.component_revision), ) .await?; - let input = decode_recursive_stream_value(method_parameters.clone(), |_, _| { - Ok(SchemaValueStream::from_host_endpoint(())) - }) - .map_err(|details| RpcError::ProtocolError { details })?; + let input_encoded_len = method_parameters.encoded_len(); + let input = decode_invocation_input(method_parameters) + .map_err(|details| RpcError::ProtocolError { details })?; let parsed_agent_id = ParsedAgentId::parse(&owned_agent_id.agent_id.agent_id, &component.metadata) .map_err(|details| RpcError::ProtocolError { details })?; @@ -1497,7 +1495,7 @@ impl Rpc for DirectWorkerInvocationRpc { let start = InvocationStart { agent_id: Some(owned_agent_id.agent_id().into()), method_name: Some(method_name.clone()), - input: Some(method_parameters.clone()), + input: None, idempotency_key: Some(idempotency_key.clone().into()), context: Some(golem_api_grpc::proto::golem::worker::InvocationContext { parent: Some(self_agent_id.clone().into()), @@ -1538,7 +1536,7 @@ impl Rpc for DirectWorkerInvocationRpc { component.revision, expected_callee_fingerprint, invocation, - method_parameters, + input_encoded_len, acceptance_committed, self.config() .limits diff --git a/golem-worker-executor/src/worker/mod.rs b/golem-worker-executor/src/worker/mod.rs index d3055e15be..4130f2bafb 100644 --- a/golem-worker-executor/src/worker/mod.rs +++ b/golem-worker-executor/src/worker/mod.rs @@ -121,9 +121,10 @@ use golem_common::model::worker::{ AgentConfigEntryDto, ResolvedRevert, RevertWorkerTarget, TypedAgentConfigEntry, }; use golem_common::model::{ - AgentFingerprint, AgentId, AgentInvocation, AgentInvocationOutput, AgentInvocationResult, - AgentMetadata, AgentStatusRecord, IdempotencyKey, OwnedAgentId, PendingInvocationRef, - PendingUpdateKind, PendingUpdateRef, Timestamp, TimestampedAgentInvocation, + AgentFingerprint, AgentId, AgentInvocation, AgentInvocationOutput, AgentInvocationPayload, + AgentInvocationResult, AgentMetadata, AgentStatusRecord, IdempotencyKey, OwnedAgentId, + PendingInvocationRef, PendingUpdateKind, PendingUpdateRef, Timestamp, + TimestampedAgentInvocation, }; use golem_common::one_shot::OneShotEvent; use golem_common::read_only_lock; @@ -580,6 +581,24 @@ impl UsesAllDeps for Worker { } } +fn into_pending_invocation_parts( + invocation: AgentInvocation, +) -> ( + Option, + IdempotencyKey, + AgentInvocationPayload, + InvocationContextStack, +) { + let semantic_idempotency_key = invocation.idempotency_key().cloned(); + let (storage_idempotency_key, payload, invocation_context) = invocation.into_parts(); + ( + semantic_idempotency_key, + storage_idempotency_key, + payload, + invocation_context, + ) +} + impl Worker { pub(crate) fn durable_stream_consumer_journal(&self) -> Arc { Arc::new(WorkerDurableStreamConsumerJournal { @@ -3126,17 +3145,17 @@ impl Worker { } } - let (idempotency_key, invocation_payload, invocation_context) = invocation.into_parts(); + let ( + semantic_idempotency_key, + idempotency_key, + invocation_payload, + invocation_context, + ) = into_pending_invocation_parts(invocation); let invocation_context = invocation_context .limit_depth(self.deps.config().limits.max_invocation_context_stack_depth); - let invocation = AgentInvocation::from_parts( - idempotency_key.clone(), - invocation_payload.clone(), - invocation_context.clone(), - ); let payload = self .oplog - .upload_payload(&invocation_payload) + .upload_payload_owned(invocation_payload) .await .map_err(|e| { WorkerExecutorError::invalid_request(format!( @@ -3151,10 +3170,6 @@ impl Worker { invocation_context.trace_states, invocation_context_spans, ); - let timestamped_invocation = TimestampedAgentInvocation { - timestamp: entry.timestamp(), - invocation, - }; // Snapshot the epoch under the instance lock that commits the // pending entry. Read-only captures the current epoch for later @@ -3176,7 +3191,7 @@ impl Worker { self.add_and_commit_oplog_internal(&instance_guard, entry, None) .await; - if let Some(idempotency_key) = timestamped_invocation.invocation.idempotency_key() { + if let Some(idempotency_key) = semantic_idempotency_key { // Captured here, inside the producer span, because a consumer links // back to the *creation context* of the work rather than to wherever // the caller happened to call from. @@ -3189,7 +3204,7 @@ impl Worker { self.external_invocation_origins .write() .await - .insert(idempotency_key.clone(), origin); + .insert(idempotency_key, origin); } if let WorkerInstance::Running(running) = &*instance_guard { @@ -3388,7 +3403,7 @@ impl Worker { .limit_depth(self.deps.config().limits.max_invocation_context_stack_depth); let payload = self .oplog - .upload_payload(&invocation_payload) + .upload_payload_owned(invocation_payload) .await .map_err(|error| { WorkerExecutorError::invalid_request(format!( @@ -7696,6 +7711,22 @@ mod tests { use std::path::Path; use test_r::test; + #[test] + fn pending_manual_update_keeps_storage_key_but_has_no_semantic_key() { + let target_revision = ComponentRevision::new(2).unwrap(); + let (semantic_key, storage_key, payload, _) = + into_pending_invocation_parts(AgentInvocation::ManualUpdate { target_revision }); + + assert!(semantic_key.is_none()); + assert!(!storage_key.value.is_empty()); + assert!(matches!( + payload, + AgentInvocationPayload::ManualUpdate { + target_revision: actual + } if actual == target_revision + )); + } + #[test] fn reconstruction_agent_quota_maps_to_startup_suspension() { let error =