diff --git a/.gitignore b/.gitignore index 592dd2057c..251e82a649 100644 --- a/.gitignore +++ b/.gitignore @@ -69,3 +69,6 @@ docs/snippets/loop.log /sdks/scala/out/ /sdks/scala/mill/out/ + +/sdks/moonbit/_build +/sdks/moonbit/.mooncakes diff --git a/Cargo.lock b/Cargo.lock index 339fc33d13..43a6a66c7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3700,6 +3700,7 @@ dependencies = [ "futures-util", "fuzzy-matcher", "goldenfile", + "golem-api-grpc", "golem-client", "golem-common", "heck", @@ -3725,6 +3726,7 @@ dependencies = [ "prettyplease", "proc-macro2", "proptest", + "prost 0.14.3", "quote", "regex", "reqwest 0.13.2", @@ -3750,6 +3752,7 @@ dependencies = [ "tokio", "tokio-stream", "tokio-tungstenite 0.25.0", + "tokio-util", "toml 0.9.12+spec-1.1.0", "toml_edit 0.23.10+spec-1.0.0", "tracing", @@ -4058,6 +4061,7 @@ dependencies = [ "base64 0.22.1", "bigdecimal", "bit-vec 0.6.3", + "blake3", "bytes", "chrono", "combine", @@ -4281,6 +4285,7 @@ dependencies = [ "applying", "arc-swap", "assert2", + "async-broadcast", "async-lock", "async-recursion", "async-scoped", @@ -4401,6 +4406,7 @@ dependencies = [ "tempfile", "test-r", "tokio", + "tokio-stream", "tokio-util", "tonic 0.14.5", "tonic-tracing-opentelemetry", @@ -5303,6 +5309,7 @@ dependencies = [ "opentelemetry 0.30.0", "opentelemetry_sdk 0.30.0", "pretty_assertions", + "prost 0.14.3", "rand 0.9.2", "redis", "reqwest 0.13.2", diff --git a/Cargo.toml b/Cargo.toml index 95676aa6ec..cbd666a56c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,7 @@ applying = "1.0.1" arc-swap = "1.9.1" arbitrary = "1.4.1" assert2 = "0.3.15" +async-broadcast = "0.7.2" async-fs = "2.1.2" async-hash = "0.5.4" async-lock = "3.4.0" diff --git a/cli/golem-cli/Cargo.toml b/cli/golem-cli/Cargo.toml index 1d863d5f9c..9f1e8be927 100644 --- a/cli/golem-cli/Cargo.toml +++ b/cli/golem-cli/Cargo.toml @@ -29,6 +29,7 @@ test = true harness = false [dependencies] +golem-api-grpc = { workspace = true } golem-client = { workspace = true } golem-common = { workspace = true, default-features = true } @@ -72,6 +73,7 @@ phf = { workspace = true } portable-pty = { workspace = true } prettyplease = { workspace = true } proc-macro2 = { workspace = true } +prost = { workspace = true } quote = { workspace = true } regex = { workspace = true } reqwest = { workspace = true } @@ -92,9 +94,10 @@ tempfile = { workspace = true } terminal_size = { workspace = true } textwrap = { workspace = true } thiserror = { workspace = true } -tokio = { workspace = true } +tokio = { workspace = true, features = ["io-util"] } tokio-stream = { workspace = true } tokio-tungstenite = { workspace = true } +tokio-util = { workspace = true } toml = { workspace = true } toml_edit = { workspace = true } tracing = { workspace = true } diff --git a/cli/golem-cli/command-output-schema/command-output.schema.json b/cli/golem-cli/command-output-schema/command-output.schema.json index 9375eface4..6a4ad9bef4 100644 --- a/cli/golem-cli/command-output-schema/command-output.schema.json +++ b/cli/golem-cli/command-output-schema/command-output.schema.json @@ -78,6 +78,9 @@ { "$ref": "#/definitions/agent.invoke" }, + { + "$ref": "#/definitions/agent.invoke-session" + }, { "$ref": "#/definitions/agent.list" }, @@ -1559,6 +1562,63 @@ }, "additionalProperties": false }, + "agent.invoke-session": { + "type": "object", + "description": "One lifecycle document emitted per accepted, result, stream, or terminal event by `golem agent invoke` in structured formats. Parse stdout as a sequence of documents, not as one array or object.", + "x-golem-output-mode": "multi-document", + "x-golem-command": "agent invoke", + "required": [ + "$type", + "kind", + "idempotencyKey" + ], + "properties": { + "$type": { + "const": "agent.invoke-session" + }, + "kind": { + "$ref": "#/definitions/AgentInvocationSessionEventKind" + }, + "idempotencyKey": { + "type": "string" + }, + "agentId": { + "type": "string" + }, + "componentRevision": { + "type": "integer", + "minimum": 0 + }, + "outcome": { + "type": "string" + }, + "reason": { + "type": "string" + }, + "error": { + "type": "string" + }, + "streamId": { + "type": "integer", + "minimum": 0 + }, + "parentStreamId": { + "type": "integer", + "minimum": 0 + }, + "path": { + "type": "string" + }, + "offset": { + "type": "integer", + "minimum": 0 + }, + "value": { + "$ref": "#/definitions/JsonValue" + } + }, + "additionalProperties": false + }, "agent.list": { "type": "object", "description": "Single structured output document emitted by `golem agent list`.", @@ -4788,6 +4848,19 @@ "missed-messages" ] }, + "AgentInvocationSessionEventKind": { + "type": "string", + "enum": [ + "accepted", + "rejected", + "result", + "item", + "end", + "stream-error", + "stream-cancel", + "finished" + ] + }, "AgentDeletionMeta": { "type": "object", "required": ["componentName", "agentId"], @@ -7373,6 +7446,10 @@ "type": "agent.invoke", "rustType": "InvokeResultView" }, + { + "type": "agent.invoke-session", + "rustType": "AgentInvocationSessionEvent" + }, { "type": "agent.list", "rustType": "AgentsMetadataResponseView" diff --git a/cli/golem-cli/src/bridge_gen/moonbit/mod.rs b/cli/golem-cli/src/bridge_gen/moonbit/mod.rs index 397bb562ea..5b09f908f7 100644 --- a/cli/golem-cli/src/bridge_gen/moonbit/mod.rs +++ b/cli/golem-cli/src/bridge_gen/moonbit/mod.rs @@ -406,6 +406,7 @@ impl MoonBitBridgeGenerator { "golemcloud/golem_sdk/interface/wasi/clocks/system-clock" @systemClock, "golemcloud/golem_sdk/rpc", "golemcloud/golem_sdk/schema_model" @model, + "golemcloud/golem_sdk/schema_model_host" @model_host, }} "#}, }; @@ -1448,7 +1449,7 @@ fn guest_decode_unstructured_binary(value : @model.SchemaValue, allowed : Array[ writer.indent(); writer.line("Some(tree) => {"); writer.indent(); - writer.line("let value = @model.schema_value_from_wit(tree) catch {"); + writer.line("let value = @model_host.schema_value_from_wit(tree) catch {"); writer.indent(); writer.line(format!( "error => raise @common.AgentError::InvalidType({} + error.to_string())", @@ -2924,6 +2925,12 @@ fn moonbit_string_literal(value: &str) -> String { escaped } +fn moonbit_char_option_literal(value: Option) -> String { + value + .map(|value| format!("Some('\\u{{{:x}}}')", value as u32)) + .unwrap_or_else(|| "None".to_string()) +} + /// Whether a (ref-resolved) schema type becomes a generated MoonBit definition /// (struct / enum). Other named defs are inlined at their use sites. fn is_named_composite(resolved: &SchemaType) -> bool { diff --git a/cli/golem-cli/src/bridge_gen/moonbit/tool.rs b/cli/golem-cli/src/bridge_gen/moonbit/tool.rs index e3e7c709d3..015bcf4c51 100644 --- a/cli/golem-cli/src/bridge_gen/moonbit/tool.rs +++ b/cli/golem-cli/src/bridge_gen/moonbit/tool.rs @@ -17,7 +17,8 @@ use super::moonbit::{ to_moonbit_constructor_ident, to_moonbit_term_ident, unique_idents_with_reserved, }; use super::{ - MoonBitBridgeGenerator, emit_schema_graph_literal, guest_codec_source, moonbit_string_literal, + MoonBitBridgeGenerator, emit_schema_graph_literal, guest_codec_source, + moonbit_char_option_literal, moonbit_string_literal, }; use crate::bridge_gen::tool_bridge_client_directory_name; use crate::bridge_gen::tool_common::{ @@ -581,9 +582,10 @@ impl MoonBitToolBridgeGenerator { .zip(&field_graphs) .map(|(field, field_graph)| { format!( - "@tool.CanonicalInputField::{{ name: {}, aliases: {}, type_: {} }}", + "@tool.CanonicalInputField::{{ name: {}, aliases: {}, short: {}, type_: {} }}", moonbit_string_literal(&field.name), moonbit_string_array(&field.aliases), + moonbit_char_option_literal(field.short), emit_schema_graph_literal(&field_graph.graph) ) }) diff --git a/cli/golem-cli/src/bridge_gen/scala/mod.rs b/cli/golem-cli/src/bridge_gen/scala/mod.rs index 4333a56b8e..fdbe7f7876 100644 --- a/cli/golem-cli/src/bridge_gen/scala/mod.rs +++ b/cli/golem-cli/src/bridge_gen/scala/mod.rs @@ -920,38 +920,43 @@ impl ScalaBridgeGenerator { )); writer.indent(); writer.line(format!( - "val parameters = {GUEST_CODEC}.encodeValue(methodParameters({invoke_args}))" + "{GUEST_CODEC}.encodeValueAsync(methodParameters({invoke_args})).flatMap {{ parameters =>" )); writer.line(format!( - "{GUEST_RUNTIME_PKG}.FutureInterop.fromEither(resolved.invokeAndAwaitWithMetadata({method_name_lit}, parameters)).map {{ __response =>" + "{GUEST_RUNTIME_PKG}.FutureInterop.fromEither(resolved.cancelableAsyncInvokeAndAwaitWithMetadata({method_name_lit}, parameters)).flatMap {{ __invocation =>" )); writer.indent(); - writer.line("val __result = __response.value"); + writer.line("__invocation.result.map { __result =>"); + writer.indent(); writer.line("val __decoded = {"); writer.indent(); writer.line(decode_block.clone()); writer.dedent(); writer.line("}"); writer.line(format!( - "{GUEST_RUNTIME_PKG}.runtime.rpc.InvocationResult(__response.metadata, __decoded)" + "{GUEST_RUNTIME_PKG}.runtime.rpc.InvocationResult(__invocation.metadata, __decoded)" )); writer.dedent(); writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); writer.dedent(); + writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); + writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); + writer.dedent(); writer.line("}"); } else { writer.line(format!("def apply({param_decls}): {FUTURE}[{ret_ty}] = {{")); writer.indent(); writer.line(format!( - "val parameters = {GUEST_CODEC}.encodeValue(methodParameters({invoke_args}))" + "{GUEST_CODEC}.encodeValueAsync(methodParameters({invoke_args})).flatMap {{ parameters =>" )); writer.line(format!( - "{GUEST_RUNTIME_PKG}.FutureInterop.fromEither(resolved.invokeAndAwait({method_name_lit}, parameters)).map {{ __result =>" + "resolved.asyncInvokeAndAwait({method_name_lit}, parameters).map {{ __result =>" )); writer.indent(); writer.line(decode_block.clone()); writer.dedent(); writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); + writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); writer.dedent(); writer.line("}"); } @@ -959,14 +964,14 @@ impl ScalaBridgeGenerator { if ephemeral { writer.line(format!( - "def cancelable({param_decls}): _root_.scala.Either[{STRING}, {GUEST_RUNTIME_PKG}.runtime.rpc.CancelableAsyncInvocation[{ret_ty}]] = {{" + "def cancelable({param_decls}): {FUTURE}[{GUEST_RUNTIME_PKG}.runtime.rpc.CancelableAsyncInvocation[{ret_ty}]] = {{" )); writer.indent(); writer.line(format!( - "val parameters = {GUEST_CODEC}.encodeValue(methodParameters({invoke_args}))" + "{GUEST_CODEC}.encodeValueAsync(methodParameters({invoke_args})).flatMap {{ parameters =>" )); writer.line(format!( - "resolved.cancelableAsyncInvokeAndAwaitWithMetadata({method_name_lit}, parameters).map {{ __invocation =>" + "{GUEST_RUNTIME_PKG}.FutureInterop.fromEither(resolved.cancelableAsyncInvokeAndAwaitWithMetadata({method_name_lit}, parameters)).map {{ __invocation =>" )); writer.indent(); writer.line("val __future = __invocation.result.map { __result =>"); @@ -978,7 +983,8 @@ impl ScalaBridgeGenerator { "{GUEST_RUNTIME_PKG}.runtime.rpc.CancelableAsyncInvocation(__invocation.metadata, __future, __invocation.cancellationToken)" )); writer.dedent(); - writer.line("}"); + writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); + writer.line("}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)"); writer.dedent(); writer.line("}"); } else { @@ -987,17 +993,16 @@ impl ScalaBridgeGenerator { )); writer.indent(); writer.line(format!( - "val parameters = {GUEST_CODEC}.encodeValue(methodParameters({invoke_args}))" + "var __underlying = _root_.scala.Option.empty[_root_.golem.runtime.rpc.CancellationToken]\nvar __cancelled = false\nval __token = _root_.golem.runtime.rpc.CancellationToken.fromFunction(() => {{ __cancelled = true; __underlying.foreach(_.cancel()) }})\nval __future = {GUEST_CODEC}.encodeValueAsync(methodParameters({invoke_args})).flatMap {{ parameters =>" )); writer.line(format!( - "val (__future, __token) = resolved.cancelableAsyncInvokeAndAwait({method_name_lit}, parameters)" + "val (__rawFuture, __rawToken) = resolved.cancelableAsyncInvokeAndAwait({method_name_lit}, parameters)\n__underlying = _root_.scala.Some(__rawToken)\nif (__cancelled) __rawToken.cancel()\n__rawFuture.map {{ __result =>" )); - writer.line("(__future.map { __result =>"); writer.indent(); writer.line(decode_block.clone()); writer.dedent(); writer.line( - "}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue), __token)", + "}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)\n}(_root_.scala.scalajs.concurrent.JSExecutionContext.Implicits.queue)\n(__future, __token)", ); writer.dedent(); writer.line("}"); @@ -3609,7 +3614,7 @@ mod tests { assert!(client_source.contains("_root_.golem.schema.SchemaValue.StringValue(message)")); assert!(client_source.contains("_root_.golem.runtime.rpc.SchemaRpcCodec.encodeValue")); assert!(client_source.contains("_root_.golem.runtime.rpc.RemoteAgentClient.resolve")); - assert!(client_source.contains("resolved.invokeAndAwait")); + assert!(client_source.contains("resolved.asyncInvokeAndAwait")); assert!(client_source.contains("def cancelable(")); assert!(client_source.contains("resolved.cancelableAsyncInvokeAndAwait")); assert!(client_source.contains("def scheduleCancelableAt(")); @@ -3657,6 +3662,44 @@ mod tests { ); } + #[test] + fn ephemeral_guest_await_remains_asynchronous() { + let dir = TempDir::new().unwrap(); + let target_path = + Utf8PathBuf::from_path_buf(dir.path().join("alpha-agent-guest-client")).unwrap(); + let mut agent_type = minimal_agent_type("AlphaAgent"); + agent_type.mode = AgentMode::Ephemeral; + agent_type.methods.push(AgentMethodSchema { + name: "echo".to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::parameters(vec![]), + output_schema: OutputSchema::Unit, + http_endpoint: vec![], + read_only: None, + }); + let mut generator = ScalaBridgeGenerator::new_with_mode( + agent_type, + &target_path, + true, + ScalaBridgeMode::GuestWasmRpc, + ) + .unwrap(); + + generator.generate().unwrap(); + + let client_source = std::fs::read_to_string( + target_path + .join("src/main/scala/golem/bridge/client/alpha_agent/AlphaAgentClient.scala"), + ) + .unwrap(); + assert!( + client_source.contains("resolved.cancelableAsyncInvokeAndAwaitWithMetadata("), + "ephemeral await must use asynchronous metadata-aware RPC:\n{client_source}" + ); + assert!(!client_source.contains("resolved.invokeAndAwaitWithMetadata(")); + } + #[test] fn guest_generation_renames_parameters_that_collide_with_cancelable_locals() { let dir = TempDir::new().unwrap(); @@ -3694,7 +3737,7 @@ mod tests { .unwrap(); assert!(client_source.contains("def cancelable(__future_2:")); assert!(client_source.contains("__token_2:")); - assert!(client_source.contains("val (__future, __token) =")); + assert!(client_source.contains("val (__rawFuture, __rawToken) =")); assert!(!client_source.contains("def cancelable(__future:")); } diff --git a/cli/golem-cli/src/command.rs b/cli/golem-cli/src/command.rs index 316ae6178d..7346dddb75 100644 --- a/cli/golem-cli/src/command.rs +++ b/cli/golem-cli/src/command.rs @@ -1291,7 +1291,7 @@ pub mod worker { }; use crate::model::agent::{AgentListMode, AgentUpdateMode}; use chrono::{DateTime, Utc}; - use clap::Subcommand; + use clap::{Subcommand, ValueEnum}; use golem_client::model::ScanCursor; use golem_common::model::IdempotencyKey; use golem_common::model::agent::AgentTypeName; @@ -1299,6 +1299,18 @@ pub mod worker { use golem_common::model::worker::AgentConfigEntryDto; use uuid::Uuid; + #[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] + pub enum InvocationStdinFormat { + Value, + Raw, + } + + #[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] + pub enum InvocationStdoutFormat { + Value, + Raw, + } + #[derive(Debug, Subcommand)] pub enum AgentSubcommand { /// Create new agent @@ -1352,6 +1364,15 @@ pub mod worker { no_stream: bool, #[command(flatten)] stream_args: StreamArgs, + /// Framing used when `-` binds stdin to a direct stream parameter. + /// `value` parses one source-language value per line; `raw` accepts + /// only stream or stream. + #[arg(long, value_enum, default_value = "value")] + stdin_format: InvocationStdinFormat, + /// Rendering used for invocation result streams. `raw` writes only + /// bytes and requires one direct stream or stream result. + #[arg(long, value_enum, default_value = "value")] + stdout_format: InvocationStdoutFormat, #[command(flatten)] post_deploy_args: Option, /// Schedule the invocation at a specific time (ISO 8601 / RFC 3339 format, e.g. 2026-03-15T10:30:00Z) diff --git a/cli/golem-cli/src/command_handler/agent/invocation_session.rs b/cli/golem-cli/src/command_handler/agent/invocation_session.rs new file mode 100644 index 0000000000..3cba2b7cb7 --- /dev/null +++ b/cli/golem-cli/src/command_handler/agent/invocation_session.rs @@ -0,0 +1,2550 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::agent_id_display::{SourceLanguage, render_schema_value}; +use crate::command::worker::{InvocationStdinFormat, InvocationStdoutFormat}; +use crate::command_handler::agent::parse_method_argument_schema_value; +use crate::command_handler::log::render_command_output_document; +use crate::context::Context; +use crate::error::{NonSuccessfulExit, PipedExitCode}; +use crate::model::agent::invocation_session::{ + AgentInvocationSessionEvent, AgentInvocationSessionEventKind, +}; +use crate::model::format::Format; +use anyhow::{Context as _, anyhow, bail}; +use futures_util::{SinkExt, StreamExt}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::schema::{ + BinaryValue, RecordValue, SchemaValue as ProtoSchemaValue, SchemaValueStreamReference, + schema_value, +}; +use golem_api_grpc::proto::golem::worker::{ + InputStreamEnd, InputStreamItem, InvocationResponse, PublicInvocationRequest, + PublicInvocationStart, StreamCancel, StreamCancelReason, StreamCancelRole, input_stream_item, + invocation_response, invocation_session_completion, invocation_session_result, + public_invocation_request, +}; +use golem_common::model::IdempotencyKey; +use golem_common::model::worker::AgentConfigEntryDto; +use golem_common::schema::agent::{ + AgentMethodSchema, AgentTypeSchema, OutputSchema, ParsedAgentId, +}; +use golem_common::schema::{BinaryValuePayload, SchemaGraph, SchemaType, SchemaValue}; +use native_tls::TlsConnector; +use prost::Message as _; +use std::collections::HashMap; +use std::io::{BufRead, BufReader, ErrorKind, Read, Write}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::{mpsc, oneshot}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::{self, Message}; +use tokio_tungstenite::{Connector, connect_async_tls_with_config}; +use tokio_util::sync::CancellationToken; + +const PIPELINE_CAPACITY: usize = 16; +const RAW_CHUNK_SIZE: usize = 64 * 1024; + +#[derive(Clone, Debug)] +struct InputBinding { + stream_id: u64, + parameter_name: String, + item_type: SchemaType, + raw_kind: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RawStreamKind { + Binary, + U8, +} + +#[derive(Clone, Debug)] +struct OutputStream { + item_type: SchemaType, + parent_stream_id: Option, + path: String, + raw_kind: Option, + next_offset: u64, + terminal: bool, +} + +enum OutputJob { + Text(String), + Raw(Vec), + Event(Box), +} + +struct OutputChannel { + tx: mpsc::Sender, + interrupt: CancellationToken, + input_failed: CancellationToken, +} + +struct InputFailure { + error: anyhow::Error, + reason: StreamCancelReason, +} + +#[derive(Debug, thiserror::Error)] +#[error("stdin input failed")] +struct InputFailureSignal; + +#[derive(Default)] +struct SessionIdentity { + agent_id: Option, + component_revision: Option, +} + +pub(super) struct InvocationSessionArgs { + pub application_name: String, + pub environment_name: String, + pub agent_type: AgentTypeSchema, + pub parsed_agent_id: ParsedAgentId, + pub method_name: String, + pub arguments: Vec, + pub config: Vec, + pub idempotency_key: IdempotencyKey, + pub stdin_format: InvocationStdinFormat, + pub stdout_format: InvocationStdoutFormat, +} + +pub(super) async fn invoke(ctx: Arc, args: InvocationSessionArgs) -> anyhow::Result<()> { + let method = args + .agent_type + .methods + .iter() + .find(|method| method.name == args.method_name) + .cloned() + .ok_or_else(|| anyhow!("Method '{}' not found in agent type", args.method_name))?; + let source_language = SourceLanguage::from(args.agent_type.source_language.clone()); + let (method_parameters, input_binding) = prepare_method_parameters( + &args.agent_type.schema, + &method, + args.arguments, + &source_language, + args.stdin_format, + )?; + if args.stdin_format == InvocationStdinFormat::Raw && input_binding.is_none() { + bail!("--stdin-format raw requires stdin bound to stream or stream with '-'"); + } + let input_stream_id = input_binding.as_ref().map(|binding| binding.stream_id); + validate_stdout_format( + &args.agent_type.schema, + &method.output_schema, + args.stdout_format, + )?; + + let constructor_parameters = args + .parsed_agent_id + .parameters + .value() + .clone() + .try_into() + .map_err(anyhow::Error::msg)?; + let idempotency_key_value = args.idempotency_key.value.clone(); + let start = PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: args.application_name, + environment_name: args.environment_name, + agent_type_name: args.parsed_agent_id.agent_type.to_string(), + constructor_parameters: Some(constructor_parameters), + phantom_id: args.parsed_agent_id.phantom_id.map(Into::into), + config: args.config.into_iter().map(Into::into).collect(), + method_name: args.method_name, + method_parameters: Some(method_parameters), + idempotency_key: Some(args.idempotency_key.into()), + }, + )), + }; + + let request = websocket_request(&ctx).await?; + let connector = if ctx.allow_insecure() { + Some(Connector::NativeTls( + TlsConnector::builder() + .danger_accept_invalid_certs(true) + .danger_accept_invalid_hostnames(true) + .build()?, + )) + } else { + None + }; + let (socket, _) = connect_async_tls_with_config(request, None, false, connector) + .await + .context("failed to connect to the agent invocation session")?; + let (mut socket_sink, mut socket_stream) = socket.split(); + + let (wire_tx, mut wire_rx) = mpsc::channel::(PIPELINE_CAPACITY); + let mut wire_writer = tokio::spawn(async move { + while let Some(message) = wire_rx.recv().await { + socket_sink.send(message).await?; + } + socket_sink.close().await + }); + + let interrupt = CancellationToken::new(); + let signal_interrupt = interrupt.clone(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + signal_interrupt.cancel(); + } + }); + + let (input_tx, mut input_rx) = mpsc::channel::(PIPELINE_CAPACITY); + let (input_failure_tx, mut input_failure_rx) = oneshot::channel::(); + let input_cancelled = CancellationToken::new(); + let input_failed = CancellationToken::new(); + let stdin_format = args.stdin_format; + let stdin_source_language = source_language.clone(); + let stdin_graph = args.agent_type.schema.clone(); + let has_input = input_binding.is_some(); + let mut input_failure_open = has_input; + if let Some(binding) = input_binding { + let reader_cancelled = input_cancelled.clone(); + let reader_failed = input_failed.clone(); + std::thread::spawn(move || { + if let Err(failure) = read_stdin( + binding, + stdin_format, + stdin_source_language, + stdin_graph, + &input_tx, + &reader_cancelled, + ) && input_failure_tx.send(failure).is_ok() + { + reader_failed.cancel(); + } + }); + } + + let format = ctx.format(); + let structured = format.is_structured() && args.stdout_format == InvocationStdoutFormat::Value; + let (output_job_tx, output_rx) = mpsc::channel::(PIPELINE_CAPACITY); + let output_tx = OutputChannel { + tx: output_job_tx, + interrupt: interrupt.clone(), + input_failed: input_failed.clone(), + }; + let (output_result_tx, mut output_result_rx) = oneshot::channel(); + let colorize = ctx.should_colorize(); + std::thread::spawn(move || { + let _ = output_result_tx.send(write_output(output_rx, format, colorize)); + }); + + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&start) + .map_err(anyhow::Error::msg)?; + send_request(&wire_tx, start, &interrupt).await?; + + let mut output_streams = HashMap::::new(); + let mut failed = false; + let mut accepted = false; + let mut stdin_open = has_input; + let mut input_terminal = !has_input; + let mut pending_input_items = 0_usize; + let mut acknowledged_input_offset = 0_u64; + let mut pending_input_request = None; + let mut fatal_input_failure = None; + let mut session_identity = SessionIdentity::default(); + let mut wire_complete = false; + + 'session: while !state.is_complete() { + tokio::select! { + biased; + _ = interrupt.cancelled() => { + input_cancelled.cancel(); + input_rx.close(); + cancel_open_streams( + &mut state, + &wire_tx, + if accepted && !input_terminal { + input_stream_id.map(|stream_id| (stream_id, acknowledged_input_offset)) + } else { + None + }, + &output_streams, + StreamCancelReason::Cancelled, + "invocation interrupted by the client", + ); + drop(wire_tx); + if !wire_complete && tokio::time::timeout(Duration::from_secs(3), &mut wire_writer).await.is_err() { + wire_writer.abort(); + } + bail!(PipedExitCode(130)); + } + failure = &mut input_failure_rx, if input_failure_open => { + input_failure_open = false; + if let Ok(failure) = failure { + fatal_input_failure = Some(failure); + break 'session; + } + } + result = &mut wire_writer, if !wire_complete => { + match result { + Ok(Ok(())) => bail!("agent invocation session connection closed before completion"), + Ok(Err(error)) if is_connection_closed(&error) => { + wire_complete = true; + input_cancelled.cancel(); + input_rx.close(); + while input_rx.try_recv().is_ok() {} + pending_input_request = None; + stdin_open = false; + } + Ok(Err(error)) => return Err(error.into()), + Err(error) => return Err(error.into()), + } + } + result = &mut output_result_rx => { + match result { + Ok(Ok(())) => bail!("invocation output closed before completion"), + Ok(Err(error)) if error.downcast_ref::().is_some_and(|exit| exit.0 == 0) => { + input_cancelled.cancel(); + input_rx.close(); + cancel_open_streams( + &mut state, + &wire_tx, + if accepted && !input_terminal { + input_stream_id.map(|stream_id| (stream_id, acknowledged_input_offset)) + } else { + None + }, + &output_streams, + StreamCancelReason::Cancelled, + "invocation output was closed by the consumer", + ); + drop(wire_tx); + if !wire_complete && tokio::time::timeout(Duration::from_secs(3), &mut wire_writer).await.is_err() { + wire_writer.abort(); + } + return Err(error); + } + Ok(Err(error)) => return Err(error), + Err(_) => bail!("invocation output writer stopped unexpectedly"), + } + } + request = input_rx.recv(), if accepted && stdin_open && pending_input_items < PIPELINE_CAPACITY && pending_input_request.is_none() => { + match request { + Some(request) => pending_input_request = Some(request), + None => stdin_open = false, + } + } + permit = wire_tx.clone().reserve_owned(), if pending_input_request.is_some() => { + let permit = permit.map_err(|_| anyhow!("agent invocation session connection closed"))?; + let request = pending_input_request + .take() + .expect("wire send selected without a pending input request"); + state.validate_public_request(&request).map_err(anyhow::Error::msg)?; + if matches!(request.request, Some(public_invocation_request::Request::InputItem(_))) { + pending_input_items += 1; + } + if matches!(request.request, Some(public_invocation_request::Request::InputEnd(_))) { + input_terminal = true; + } + permit.send(encode_request(request)?); + } + frame = socket_stream.next() => { + let response = match receive_response( + frame, + &wire_tx, + &interrupt, + &input_failed, + ).await { + Ok(Some(response)) => response, + Ok(None) => continue, + Err(error) if error.downcast_ref::().is_some() => { + input_failure_open = false; + fatal_input_failure = Some(take_input_failure(&mut input_failure_rx)?); + break 'session; + } + Err(error) => return Err(error), + }; + state.validate_response(&response).map_err(anyhow::Error::msg)?; + if matches!(response.response, Some(invocation_response::Response::Accepted(_))) { + accepted = true; + } + if let Some(invocation_response::Response::InputAck(ack)) = response.response.as_ref() { + pending_input_items = pending_input_items.checked_sub(1).ok_or_else(|| anyhow!("received an input acknowledgement without a pending item"))?; + acknowledged_input_offset = ack.sequence.checked_add(ack.logical_item_count).ok_or_else(|| anyhow!("input acknowledgement offset overflow"))?; + } + if response_cancels_input(&response, input_stream_id) { + input_cancelled.cancel(); + input_rx.close(); + while input_rx.try_recv().is_ok() {} + pending_input_request = None; + input_failure_open = false; + input_failure_rx.close(); + stdin_open = false; + input_terminal = true; + } + match handle_response( + response, + &args.agent_type.schema, + &method.output_schema, + &source_language, + args.stdout_format, + structured, + &idempotency_key_value, + &mut session_identity, + &mut output_streams, + &output_tx, + ).await { + Ok(response_failed) => failed |= response_failed, + Err(error) if error.downcast_ref::().is_some() => { + input_failure_open = false; + fatal_input_failure = Some(take_input_failure(&mut input_failure_rx)?); + break 'session; + } + Err(error) => return Err(error), + } + } + } + } + + if let Some(failure) = fatal_input_failure { + input_cancelled.cancel(); + input_rx.close(); + while input_rx.try_recv().is_ok() {} + cancel_open_streams( + &mut state, + &wire_tx, + if accepted && !input_terminal { + input_stream_id.map(|stream_id| (stream_id, acknowledged_input_offset)) + } else { + None + }, + &output_streams, + failure.reason, + &failure.error.to_string(), + ); + drop(wire_tx); + if !wire_complete + && tokio::time::timeout(Duration::from_secs(3), &mut wire_writer) + .await + .is_err() + { + wire_writer.abort(); + } + return Err(failure.error); + } + + if let Err(error) = await_clean_close( + &mut socket_stream, + &wire_tx, + &interrupt, + &mut output_result_rx, + ) + .await + { + input_cancelled.cancel(); + input_rx.close(); + cancel_open_streams( + &mut state, + &wire_tx, + if accepted && !input_terminal { + input_stream_id.map(|stream_id| (stream_id, acknowledged_input_offset)) + } else { + None + }, + &output_streams, + StreamCancelReason::Cancelled, + "invocation session stopped while waiting for the server to close", + ); + drop(wire_tx); + if !wire_complete + && tokio::time::timeout(Duration::from_secs(3), &mut wire_writer) + .await + .is_err() + { + wire_writer.abort(); + } + return Err(error); + } + + input_cancelled.cancel(); + drop(wire_tx); + drop(output_tx); + input_rx.close(); + + let mut output_complete = false; + while !wire_complete || !output_complete { + tokio::select! { + biased; + _ = interrupt.cancelled() => { + wire_writer.abort(); + bail!(PipedExitCode(130)); + } + result = &mut output_result_rx, if !output_complete => { + match result { + Ok(result) => result?, + Err(_) => bail!("invocation output writer stopped unexpectedly"), + } + output_complete = true; + } + result = &mut wire_writer, if !wire_complete => { + match result { + Ok(Ok(())) => {} + Ok(Err(error)) if state.is_complete() && is_connection_closed(&error) => {} + Ok(Err(error)) => return Err(error.into()), + Err(error) => return Err(error.into()), + } + wire_complete = true; + } + } + } + + if input_failure_open && let Ok(failure) = input_failure_rx.try_recv() { + return Err(failure.error); + } + if failed { + bail!(NonSuccessfulExit); + } + Ok(()) +} + +async fn websocket_request(ctx: &Context) -> anyhow::Result> { + let mut url = ctx.worker_service_url().clone(); + let websocket_scheme = match url.scheme() { + "http" => "ws", + "https" => "wss", + scheme => bail!("unsupported service URL scheme '{scheme}'"), + }; + url.set_scheme(websocket_scheme) + .map_err(|_| anyhow!("failed to derive WebSocket service URL"))?; + url.set_path("/v1/agents/invoke-agent-session"); + url.set_query(None); + url.set_fragment(None); + let mut request = url.as_str().into_client_request()?; + let token = ctx.auth_token().await?; + request.headers_mut().insert( + tungstenite::http::header::AUTHORIZATION, + format!("Bearer {}", token.secret()) + .parse() + .context("invalid authorization token")?, + ); + Ok(request) +} + +fn prepare_method_parameters( + graph: &SchemaGraph, + method: &AgentMethodSchema, + arguments: Vec, + source_language: &SourceLanguage, + stdin_format: InvocationStdinFormat, +) -> anyhow::Result<(ProtoSchemaValue, Option)> { + let fields = method.input_schema.fields(); + if fields.len() != arguments.len() { + bail!( + "wrong number of parameters: expected {}, got {}", + fields.len(), + arguments.len() + ); + } + + let mut values = Vec::with_capacity(fields.len()); + let mut input_binding = None; + let mut next_stream_id = 1_u64; + for (field, argument) in fields.iter().zip(arguments) { + let resolved = graph + .resolve_ref(&field.schema) + .map_err(|error| anyhow!(error.to_string()))?; + if argument == "-" { + let SchemaType::Stream { inner, .. } = resolved else { + bail!( + "stdin marker '-' can only be used for a direct stream parameter; '{}' is not a stream", + field.name + ); + }; + if input_binding.is_some() { + bail!("stdin can only be bound to one stream parameter"); + } + let item_type = inner + .as_deref() + .cloned() + .ok_or_else(|| anyhow!("untyped streams cannot be bound to stdin"))?; + if golem_common::schema::agent::contains_stream_in_graph(graph, &item_type) { + bail!( + "stdin stream parameter '{}' cannot contain nested streams", + field.name + ); + } + let raw_kind = match stdin_format { + InvocationStdinFormat::Value => None, + InvocationStdinFormat::Raw => { + Some(raw_stream_kind(graph, &item_type).ok_or_else(|| { + anyhow!("--stdin-format raw requires stream or stream") + })?) + } + }; + let stream_id = next_stream_id; + next_stream_id += 2; + values.push(ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id }, + )), + }); + input_binding = Some(InputBinding { + stream_id, + parameter_name: field.name.clone(), + item_type, + raw_kind, + }); + } else { + if golem_common::schema::agent::contains_stream_in_graph(graph, &field.schema) { + bail!( + "stream parameter '{}' must be a direct stream bound to stdin with '-'", + field.name + ); + } + let parsed = parse_method_argument_schema_value( + &argument, + graph, + &field.schema, + source_language, + ) + .map_err(|error| anyhow!("invalid value for '{}': {}", field.name, error.message))?; + values.push(parsed.try_into().map_err(anyhow::Error::msg)?); + } + } + + Ok(( + ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: values, + })), + }, + input_binding, + )) +} + +fn validate_stdout_format( + graph: &SchemaGraph, + output: &OutputSchema, + format: InvocationStdoutFormat, +) -> anyhow::Result<()> { + if format == InvocationStdoutFormat::Value { + return Ok(()); + } + let ty = output + .schema() + .ok_or_else(|| anyhow!("--stdout-format raw requires a stream result"))?; + let resolved = graph + .resolve_ref(ty) + .map_err(|error| anyhow!(error.to_string()))?; + let SchemaType::Stream { + inner: Some(inner), .. + } = resolved + else { + bail!("--stdout-format raw requires one direct stream result"); + }; + if raw_stream_kind(graph, inner).is_none() { + bail!("--stdout-format raw requires stream or stream"); + } + Ok(()) +} + +fn raw_stream_kind(graph: &SchemaGraph, ty: &SchemaType) -> Option { + match graph.resolve_ref(ty).ok()? { + SchemaType::Binary { .. } => Some(RawStreamKind::Binary), + SchemaType::U8 { .. } => Some(RawStreamKind::U8), + _ => None, + } +} + +fn read_stdin( + binding: InputBinding, + format: InvocationStdinFormat, + source_language: SourceLanguage, + graph: SchemaGraph, + tx: &mpsc::Sender, + cancelled: &CancellationToken, +) -> Result<(), InputFailure> { + let mut offset = 0_u64; + match format { + InvocationStdinFormat::Value => { + let mut stdin = BufReader::new(std::io::stdin().lock()); + let mut line = String::new(); + loop { + if cancelled.is_cancelled() { + return Ok(()); + } + line.clear(); + match stdin.read_line(&mut line) { + Ok(0) => break, + Ok(_) => { + if line.ends_with('\n') { + line.pop(); + if line.ends_with('\r') { + line.pop(); + } + } + } + Err(error) => { + return Err(InputFailure { + error: anyhow!( + "failed to read stdin for parameter '{}' at line {}: {error}", + binding.parameter_name, + offset + 1 + ), + reason: StreamCancelReason::Transport, + }); + } + } + if cancelled.is_cancelled() { + return Ok(()); + } + let value = match parse_method_argument_schema_value( + &line, + &graph, + &binding.item_type, + &source_language, + ) { + Ok(value) => value, + Err(parse_error) => { + return Err(InputFailure { + error: anyhow!( + "invalid stdin value for parameter '{}' at line {}: {parse_error}", + binding.parameter_name, + offset + 1 + ), + reason: StreamCancelReason::Protocol, + }); + } + }; + let request = + input_value_request(binding.stream_id, offset, value).map_err(|error| { + InputFailure { + error, + reason: StreamCancelReason::Protocol, + } + })?; + if tx.blocking_send(request).is_err() { + if cancelled.is_cancelled() { + return Ok(()); + } + return Err(InputFailure { + error: anyhow!("invocation session ended while reading stdin"), + reason: StreamCancelReason::Transport, + }); + } + offset = offset.checked_add(1).ok_or_else(|| InputFailure { + error: anyhow!("stdin stream offset overflow"), + reason: StreamCancelReason::Protocol, + })?; + } + } + InvocationStdinFormat::Raw => { + let mut stdin = std::io::stdin().lock(); + let mut buffer = vec![0_u8; RAW_CHUNK_SIZE]; + loop { + if cancelled.is_cancelled() { + return Ok(()); + } + let mut count = 0; + let mut eof = false; + while count < RAW_CHUNK_SIZE { + if cancelled.is_cancelled() { + return Ok(()); + } + match stdin.read(&mut buffer[count..]) { + Ok(0) => { + eof = true; + break; + } + Ok(read) => { + count += read; + if cancelled.is_cancelled() { + return Ok(()); + } + } + Err(error) => { + return Err(InputFailure { + error: anyhow!( + "failed to read raw stdin for parameter '{}': {error}", + binding.parameter_name + ), + reason: StreamCancelReason::Transport, + }); + } + } + } + if count == 0 { + break; + } + let payload = match binding.raw_kind { + Some(RawStreamKind::Binary) => input_stream_item::Payload::Value( + SchemaValue::Binary(BinaryValuePayload { + bytes: buffer[..count].to_vec(), + mime_type: None, + }) + .try_into() + .map_err(|error| InputFailure { + error: anyhow::Error::msg(error), + reason: StreamCancelReason::Protocol, + })?, + ), + Some(RawStreamKind::U8) => { + input_stream_item::Payload::PackedU8(buffer[..count].to_vec()) + } + None => unreachable!("raw stdin was validated before reading"), + }; + if tx + .blocking_send(PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id: binding.stream_id, + sequence: offset, + payload: Some(payload), + }, + )), + }) + .is_err() + { + if cancelled.is_cancelled() { + return Ok(()); + } + return Err(InputFailure { + error: anyhow!("invocation session ended while reading stdin"), + reason: StreamCancelReason::Transport, + }); + } + let logical_item_count = if binding.raw_kind == Some(RawStreamKind::U8) { + count as u64 + } else { + 1 + }; + offset = offset + .checked_add(logical_item_count) + .ok_or_else(|| InputFailure { + error: anyhow!("stdin stream offset overflow"), + reason: StreamCancelReason::Protocol, + })?; + if eof { + break; + } + } + } + } + if cancelled.is_cancelled() { + return Ok(()); + } + if tx + .blocking_send(PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputEnd( + InputStreamEnd { + stream_id: binding.stream_id, + offset, + }, + )), + }) + .is_err() + { + if cancelled.is_cancelled() { + return Ok(()); + } + return Err(InputFailure { + error: anyhow!("invocation session ended before stdin reached EOF"), + reason: StreamCancelReason::Transport, + }); + } + Ok(()) +} + +fn input_cancel_request( + stream_id: u64, + offset: u64, + reason: StreamCancelReason, + details: String, +) -> PublicInvocationRequest { + PublicInvocationRequest { + request: Some(public_invocation_request::Request::StreamCancel( + StreamCancel { + stream_id, + offset, + role: StreamCancelRole::InputProducer as i32, + reason: reason as i32, + details: Some(details), + }, + )), + } +} + +fn output_cancel_request( + stream_id: u64, + offset: u64, + reason: StreamCancelReason, + details: String, +) -> PublicInvocationRequest { + PublicInvocationRequest { + request: Some(public_invocation_request::Request::StreamCancel( + StreamCancel { + stream_id, + offset, + role: StreamCancelRole::OutputConsumer as i32, + reason: reason as i32, + details: Some(details), + }, + )), + } +} + +fn cancel_open_streams( + state: &mut InvocationSessionState, + wire_tx: &mpsc::Sender, + input: Option<(u64, u64)>, + output_streams: &HashMap, + reason: StreamCancelReason, + details: &str, +) { + if let Some((stream_id, offset)) = input { + let request = input_cancel_request(stream_id, offset, reason, details.to_string()); + if state.validate_public_request(&request).is_ok() { + try_send_request(wire_tx, request); + } + } + + let mut open_outputs = output_streams + .iter() + .filter_map(|(stream_id, stream)| { + (!stream.terminal).then_some((*stream_id, stream.next_offset)) + }) + .collect::>(); + open_outputs.sort_unstable_by_key(|(stream_id, _)| *stream_id); + for (stream_id, offset) in open_outputs { + let request = output_cancel_request(stream_id, offset, reason, details.to_string()); + if state.validate_public_request(&request).is_ok() { + try_send_request(wire_tx, request); + } + } +} + +fn input_value_request( + stream_id: u64, + sequence: u64, + value: SchemaValue, +) -> anyhow::Result { + Ok(PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id, + sequence, + payload: Some(input_stream_item::Payload::Value( + value.try_into().map_err(anyhow::Error::msg)?, + )), + }, + )), + }) +} + +async fn send_request( + tx: &mpsc::Sender, + request: PublicInvocationRequest, + interrupt: &CancellationToken, +) -> anyhow::Result<()> { + send_message(tx, encode_request(request)?, interrupt).await +} + +fn try_send_request(tx: &mpsc::Sender, request: PublicInvocationRequest) { + if let Ok(message) = encode_request(request) { + let _ = tx.try_send(message); + } +} + +fn take_input_failure( + receiver: &mut oneshot::Receiver, +) -> anyhow::Result { + receiver.try_recv().map_err(|error| match error { + oneshot::error::TryRecvError::Empty => { + anyhow!("stdin failure signal arrived before its diagnostic") + } + oneshot::error::TryRecvError::Closed => { + anyhow!("stdin failure signal arrived without a diagnostic") + } + }) +} + +fn encode_request(request: PublicInvocationRequest) -> anyhow::Result { + let mut bytes = Vec::new(); + request.encode(&mut bytes)?; + Ok(Message::Binary(bytes.into())) +} + +async fn send_message( + tx: &mpsc::Sender, + message: Message, + interrupt: &CancellationToken, +) -> anyhow::Result<()> { + tokio::select! { + biased; + _ = interrupt.cancelled() => bail!(PipedExitCode(130)), + result = tx.send(message) => result.map_err(|_| anyhow!("agent invocation session connection closed")), + } +} + +async fn send_active_message( + tx: &mpsc::Sender, + message: Message, + interrupt: &CancellationToken, + input_failed: &CancellationToken, +) -> anyhow::Result<()> { + tokio::select! { + biased; + _ = input_failed.cancelled() => bail!(InputFailureSignal), + _ = interrupt.cancelled() => bail!(PipedExitCode(130)), + result = tx.send(message) => result.map_err(|_| anyhow!("agent invocation session connection closed")), + } +} + +async fn receive_response( + frame: Option>, + wire_tx: &mpsc::Sender, + interrupt: &CancellationToken, + input_failed: &CancellationToken, +) -> anyhow::Result> { + match frame { + Some(Ok(Message::Binary(bytes))) => Ok(Some(InvocationResponse::decode(bytes.as_slice())?)), + Some(Ok(Message::Ping(payload))) => { + send_active_message(wire_tx, Message::Pong(payload), interrupt, input_failed).await?; + Ok(None) + } + Some(Ok(Message::Pong(_))) => Ok(None), + Some(Ok(Message::Close(close))) => { + bail!("agent invocation session closed before completion: {close:?}"); + } + Some(Ok(message)) => { + bail!("unexpected WebSocket frame in invocation session: {message:?}"); + } + Some(Err(error)) => Err(error.into()), + None => bail!("agent invocation session ended before completion"), + } +} + +fn response_cancels_input(response: &InvocationResponse, input_stream_id: Option) -> bool { + let Some(input_stream_id) = input_stream_id else { + return false; + }; + matches!( + response.response.as_ref(), + Some(invocation_response::Response::StreamCancel(cancel)) + if cancel.stream_id == input_stream_id + && cancel.role() == StreamCancelRole::InputConsumer + ) +} + +async fn await_clean_close( + socket_stream: &mut S, + wire_tx: &mpsc::Sender, + interrupt: &CancellationToken, + output_result_rx: &mut oneshot::Receiver>, +) -> anyhow::Result<()> +where + S: futures_util::Stream> + Unpin, +{ + let close_timeout = tokio::time::sleep(Duration::from_secs(3)); + tokio::pin!(close_timeout); + loop { + tokio::select! { + biased; + _ = interrupt.cancelled() => bail!(PipedExitCode(130)), + result = &mut *output_result_rx => { + match result { + Ok(result) => return result, + Err(_) => bail!("invocation output writer stopped unexpectedly"), + } + } + _ = &mut close_timeout => { + bail!("invocation session did not close after completion"); + } + frame = socket_stream.next() => { + match frame { + Some(Ok(Message::Close(_))) | None => return Ok(()), + Some(Ok(Message::Ping(payload))) => { + send_message(wire_tx, Message::Pong(payload), interrupt).await?; + } + Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Binary(_))) => { + bail!("invocation session sent an event after completion"); + } + Some(Ok(message)) => { + bail!("unexpected WebSocket frame after invocation completion: {message:?}"); + } + Some(Err(error)) if is_connection_closed(&error) => return Ok(()), + Some(Err(error)) => return Err(error.into()), + } + } + } + } +} + +#[allow(clippy::too_many_arguments)] +async fn handle_response( + response: InvocationResponse, + graph: &SchemaGraph, + output_schema: &OutputSchema, + source_language: &SourceLanguage, + stdout_format: InvocationStdoutFormat, + structured: bool, + idempotency_key: &str, + session_identity: &mut SessionIdentity, + output_streams: &mut HashMap, + output_tx: &OutputChannel, +) -> anyhow::Result { + let mut failed = false; + match response.response { + Some(invocation_response::Response::Accepted(accepted)) => { + session_identity.agent_id = accepted.agent_id.as_ref().map(|id| id.name.clone()); + session_identity.component_revision = accepted.component_revision; + if structured { + let mut event = event(AgentInvocationSessionEventKind::Accepted, idempotency_key); + event.agent_id = session_identity.agent_id.clone(); + event.component_revision = session_identity.component_revision; + emit(output_tx, OutputJob::Event(event)).await?; + } + } + Some(invocation_response::Response::Rejected(rejected)) => { + failed = true; + let reason = format!("{:?}", rejected.reason()); + if structured { + let mut event = event(AgentInvocationSessionEventKind::Rejected, idempotency_key); + event.agent_id = rejected.agent_id.map(|id| id.name); + event.component_revision = rejected.component_revision; + event.reason = Some(reason); + event.error = Some(rejected.error); + emit(output_tx, OutputJob::Event(event)).await?; + } else { + eprintln!("Invocation rejected ({reason}): {}", rejected.error); + } + } + Some(invocation_response::Response::Result(result)) => { + let result_agent_id = result + .agent_id + .map(|id| id.name) + .or_else(|| session_identity.agent_id.clone()); + let result_component_revision = result + .component_revision + .or(session_identity.component_revision); + match result.result { + Some(invocation_session_result::Result::MethodResult(value)) => { + let Some(output_type) = output_schema.schema() else { + if structured { + let mut event = + event(AgentInvocationSessionEventKind::Result, idempotency_key); + event.agent_id = result_agent_id; + event.component_revision = result_component_revision; + emit(output_tx, OutputJob::Event(event)).await?; + } else { + emit(output_tx, OutputJob::Text("void".to_string())).await?; + } + return Ok(failed); + }; + discover_streams( + graph, + output_type, + &value, + "$", + None, + stdout_format, + output_streams, + )?; + if structured { + let mut event = + event(AgentInvocationSessionEventKind::Result, idempotency_key); + event.agent_id = result_agent_id; + event.component_revision = result_component_revision; + event.value = Some(proto_value_to_json(graph, output_type, &value)?); + emit(output_tx, OutputJob::Event(event)).await?; + } else if stdout_format == InvocationStdoutFormat::Value { + for (path, rendered) in + render_text_fragments(graph, output_type, &value, "$", source_language)? + { + let rendered = if path == "$" { + rendered + } else { + format!("{path}: {rendered}") + }; + emit(output_tx, OutputJob::Text(rendered)).await?; + } + } + } + Some(invocation_session_result::Result::NoResult(_)) => { + if !matches!(output_schema, OutputSchema::Unit) { + bail!("session returned no result for a value-returning method"); + } + if structured { + let mut event = + event(AgentInvocationSessionEventKind::Result, idempotency_key); + event.agent_id = result_agent_id; + event.component_revision = result_component_revision; + emit(output_tx, OutputJob::Event(event)).await?; + } else { + emit(output_tx, OutputJob::Text("void".to_string())).await?; + } + } + None => bail!("invocation result has no value"), + } + } + Some(invocation_response::Response::OutputItem(item)) => { + let stream = output_streams + .get(&item.stream_id) + .cloned() + .ok_or_else(|| anyhow!("output stream {} has no schema", item.stream_id))?; + let value = item + .value + .ok_or_else(|| anyhow!("output stream item has no value"))?; + discover_streams( + graph, + &stream.item_type, + &value, + &format!("{}[{}]", stream.path, item.offset), + Some(item.stream_id), + InvocationStdoutFormat::Value, + output_streams, + )?; + output_streams + .get_mut(&item.stream_id) + .expect("output stream disappeared while processing an item") + .next_offset = item + .offset + .checked_add(1) + .ok_or_else(|| anyhow!("output stream offset overflow"))?; + if structured { + let mut event = event(AgentInvocationSessionEventKind::Item, idempotency_key); + event.stream_id = Some(item.stream_id); + event.parent_stream_id = stream.parent_stream_id; + event.path = Some(stream.path); + event.offset = Some(item.offset); + event.value = Some(proto_value_to_json(graph, &stream.item_type, &value)?); + emit(output_tx, OutputJob::Event(event)).await?; + } else if stdout_format == InvocationStdoutFormat::Raw { + emit(output_tx, OutputJob::Raw(raw_output_bytes(&stream, value)?)).await?; + } else { + for (path, rendered) in render_text_fragments( + graph, + &stream.item_type, + &value, + &stream.path, + source_language, + )? { + let rendered = + if path != stream.path || output_streams.len() > 1 || stream.path != "$" { + format!("{path}: {rendered}") + } else { + rendered + }; + emit(output_tx, OutputJob::Text(rendered)).await?; + } + } + } + Some(invocation_response::Response::OutputEnd(end)) => { + output_streams + .get_mut(&end.stream_id) + .ok_or_else(|| anyhow!("output stream {} has no schema", end.stream_id))? + .terminal = true; + if structured { + let stream = output_streams + .get(&end.stream_id) + .ok_or_else(|| anyhow!("output stream {} has no schema", end.stream_id))?; + let mut event = event(AgentInvocationSessionEventKind::End, idempotency_key); + event.stream_id = Some(end.stream_id); + event.parent_stream_id = stream.parent_stream_id; + event.path = Some(stream.path.clone()); + event.offset = Some(end.offset); + emit(output_tx, OutputJob::Event(event)).await?; + } + } + Some(invocation_response::Response::OutputError(error)) => { + failed = true; + output_streams + .get_mut(&error.stream_id) + .ok_or_else(|| anyhow!("output stream {} has no schema", error.stream_id))? + .terminal = true; + if structured { + let stream = output_streams + .get(&error.stream_id) + .ok_or_else(|| anyhow!("output stream {} has no schema", error.stream_id))?; + let mut event = event( + AgentInvocationSessionEventKind::StreamError, + idempotency_key, + ); + event.stream_id = Some(error.stream_id); + event.parent_stream_id = stream.parent_stream_id; + event.path = Some(stream.path.clone()); + event.offset = Some(error.offset); + event.error = Some(error.details); + emit(output_tx, OutputJob::Event(event)).await?; + } else { + eprintln!( + "Output stream {} failed: {}", + error.stream_id, error.details + ); + } + } + Some(invocation_response::Response::InputAck(_)) => {} + Some(invocation_response::Response::StreamCancel(cancel)) => { + failed |= cancel.reason() != StreamCancelReason::Cancelled; + let output_stream = if cancel.role() == StreamCancelRole::OutputProducer { + let stream = output_streams + .get_mut(&cancel.stream_id) + .ok_or_else(|| anyhow!("output stream {} has no schema", cancel.stream_id))?; + stream.terminal = true; + Some(stream.clone()) + } else { + None + }; + if structured { + let mut event = event( + AgentInvocationSessionEventKind::StreamCancel, + idempotency_key, + ); + event.stream_id = Some(cancel.stream_id); + if let Some(stream) = output_stream { + event.parent_stream_id = stream.parent_stream_id; + event.path = Some(stream.path); + } + event.offset = Some(cancel.offset); + event.reason = Some(format!("{:?}", cancel.reason())); + event.error = cancel.details; + emit(output_tx, OutputJob::Event(event)).await?; + } else { + eprintln!("Stream {} was cancelled", cancel.stream_id); + } + } + Some(invocation_response::Response::AttachmentRevoked(revoked)) => { + bail!("invocation attachment was revoked: {}", revoked.details); + } + Some(invocation_response::Response::Finished(finished)) => { + let (outcome, error) = match finished.outcome { + Some(invocation_session_completion::Outcome::Success(_)) => { + ("success".to_string(), None) + } + Some(invocation_session_completion::Outcome::Failure(failure)) => { + failed = true; + (format!("{:?}", failure.kind()), Some(failure.message)) + } + None => bail!("invocation completion has no outcome"), + }; + if structured { + let mut event = event(AgentInvocationSessionEventKind::Finished, idempotency_key); + event.agent_id = session_identity.agent_id.clone(); + event.component_revision = session_identity.component_revision; + event.outcome = Some(outcome); + event.error = error; + emit(output_tx, OutputJob::Event(event)).await?; + } else if let Some(error) = error { + eprintln!("Invocation failed: {error}"); + } + } + None => bail!("empty invocation response"), + } + Ok(failed) +} + +fn event( + kind: AgentInvocationSessionEventKind, + idempotency_key: &str, +) -> Box { + Box::new(AgentInvocationSessionEvent::new(kind, idempotency_key)) +} + +async fn emit(output: &OutputChannel, job: OutputJob) -> anyhow::Result<()> { + tokio::select! { + biased; + _ = output.input_failed.cancelled() => bail!(InputFailureSignal), + _ = output.interrupt.cancelled() => bail!(PipedExitCode(130)), + result = output.tx.send(job) => result.map_err(|_| anyhow!("invocation output closed unexpectedly")), + } +} + +fn write_output( + mut rx: mpsc::Receiver, + format: Format, + colorize: bool, +) -> anyhow::Result<()> { + while let Some(job) = rx.blocking_recv() { + let bytes = render_output_job(job, format, colorize)?; + let stdout = std::io::stdout(); + write_and_flush(&mut stdout.lock(), &bytes)?; + } + Ok(()) +} + +#[cfg(test)] +fn write_output_to( + mut rx: mpsc::Receiver, + format: Format, + colorize: bool, + output: &mut W, +) -> anyhow::Result<()> { + while let Some(job) = rx.blocking_recv() { + let bytes = render_output_job(job, format, colorize)?; + write_and_flush(output, &bytes)?; + } + Ok(()) +} + +fn render_output_job(job: OutputJob, format: Format, colorize: bool) -> anyhow::Result> { + match job { + OutputJob::Text(mut text) => { + text.push('\n'); + Ok(text.into_bytes()) + } + OutputJob::Raw(bytes) => Ok(bytes), + OutputJob::Event(event) => { + let mut document = render_command_output_document(format, colorize, *event)?; + document.push('\n'); + Ok(document.into_bytes()) + } + } +} + +fn write_and_flush(output: &mut impl Write, bytes: &[u8]) -> anyhow::Result<()> { + if let Err(error) = output.write_all(bytes).and_then(|()| output.flush()) { + if error.kind() == ErrorKind::BrokenPipe { + bail!(PipedExitCode(0)); + } + return Err(error.into()); + } + Ok(()) +} + +fn raw_output_bytes(stream: &OutputStream, value: ProtoSchemaValue) -> anyhow::Result> { + match (stream.raw_kind, value.value) { + ( + Some(RawStreamKind::Binary), + Some(schema_value::Value::BinaryValue(BinaryValue { bytes, .. })), + ) => Ok(bytes), + (Some(RawStreamKind::U8), Some(schema_value::Value::U8Value(value))) => Ok(vec![ + u8::try_from(value).map_err(|_| anyhow!("u8 stream item is out of range"))?, + ]), + _ => bail!("raw output stream item does not match its declared type"), + } +} + +fn render_value( + graph: &SchemaGraph, + ty: &SchemaType, + value: ProtoSchemaValue, + source_language: &SourceLanguage, +) -> anyhow::Result { + let value = SchemaValue::try_from(value).map_err(anyhow::Error::msg)?; + Ok(render_schema_value(graph, ty, &value, source_language)) +} + +fn render_text_fragments( + graph: &SchemaGraph, + ty: &SchemaType, + value: &ProtoSchemaValue, + path: &str, + source_language: &SourceLanguage, +) -> anyhow::Result> { + let mut fragments = Vec::new(); + collect_text_fragments(graph, ty, value, path, source_language, &mut fragments)?; + Ok(fragments) +} + +fn collect_text_fragments( + graph: &SchemaGraph, + ty: &SchemaType, + value: &ProtoSchemaValue, + path: &str, + source_language: &SourceLanguage, + fragments: &mut Vec<(String, String)>, +) -> anyhow::Result<()> { + let ty = graph + .resolve_ref(ty) + .map_err(|error| anyhow!(error.to_string()))?; + let value_body = value + .value + .as_ref() + .ok_or_else(|| anyhow!("schema value at {path} is empty"))?; + + if matches!(ty, SchemaType::Stream { .. }) { + if matches!(value_body, schema_value::Value::StreamReference(_)) { + return Ok(()); + } + bail!("stream value at {path} is not a stream reference"); + } + + if let Ok(rendered) = render_value(graph, ty, value.clone(), source_language) { + fragments.push((path.to_string(), rendered)); + return Ok(()); + } + + match (ty, value_body) { + (SchemaType::Record { fields, .. }, schema_value::Value::RecordValue(record)) => { + if fields.len() != record.fields.len() { + bail!("record value at {path} has the wrong number of fields"); + } + for (field, value) in fields.iter().zip(&record.fields) { + collect_text_fragments( + graph, + &field.body, + value, + &format!("{path}.{}", field.name), + source_language, + fragments, + )?; + } + } + (SchemaType::Tuple { elements, .. }, schema_value::Value::TupleValue(tuple)) => { + if elements.len() != tuple.elements.len() { + bail!("tuple value at {path} has the wrong number of elements"); + } + for (index, (ty, value)) in elements.iter().zip(&tuple.elements).enumerate() { + collect_text_fragments( + graph, + ty, + value, + &format!("{path}[{index}]"), + source_language, + fragments, + )?; + } + } + (SchemaType::List { element, .. }, schema_value::Value::ListValue(list)) => { + for (index, value) in list.elements.iter().enumerate() { + collect_text_fragments( + graph, + element, + value, + &format!("{path}[{index}]"), + source_language, + fragments, + )?; + } + } + (SchemaType::FixedList { element, .. }, schema_value::Value::FixedListValue(list)) => { + for (index, value) in list.elements.iter().enumerate() { + collect_text_fragments( + graph, + element, + value, + &format!("{path}[{index}]"), + source_language, + fragments, + )?; + } + } + ( + SchemaType::Map { + key, + value: value_type, + .. + }, + schema_value::Value::MapValue(map), + ) => { + for (index, entry) in map.entries.iter().enumerate() { + collect_text_fragments( + graph, + key, + entry + .key + .as_ref() + .ok_or_else(|| anyhow!("map entry at {path}[{index}] has no key"))?, + &format!("{path}[{index}].key"), + source_language, + fragments, + )?; + collect_text_fragments( + graph, + value_type, + entry + .value + .as_ref() + .ok_or_else(|| anyhow!("map entry at {path}[{index}] has no value"))?, + &format!("{path}[{index}].value"), + source_language, + fragments, + )?; + } + } + (SchemaType::Option { inner, .. }, schema_value::Value::OptionValue(option)) => { + if let Some(value) = option.inner.as_deref() { + collect_text_fragments( + graph, + inner, + value, + &format!("{path}.some"), + source_language, + fragments, + )?; + } + } + (SchemaType::Variant { cases, .. }, schema_value::Value::VariantValue(variant)) => { + let case = cases + .get(variant.case as usize) + .ok_or_else(|| anyhow!("variant case at {path} is out of range"))?; + if let (Some(ty), Some(value)) = (&case.payload, variant.payload.as_deref()) { + collect_text_fragments( + graph, + ty, + value, + &format!("{path}.{}", case.name), + source_language, + fragments, + )?; + } + } + (SchemaType::Result { spec, .. }, schema_value::Value::ResultValue(result)) => { + use golem_api_grpc::proto::golem::schema::result_value::Result; + match result.result.as_ref() { + Some(Result::Ok(value)) => collect_text_fragments( + graph, + spec.ok + .as_deref() + .ok_or_else(|| anyhow!("unexpected ok payload at {path}"))?, + value, + &format!("{path}.ok"), + source_language, + fragments, + )?, + Some(Result::Err(value)) => collect_text_fragments( + graph, + spec.err + .as_deref() + .ok_or_else(|| anyhow!("unexpected err payload at {path}"))?, + value, + &format!("{path}.err"), + source_language, + fragments, + )?, + Some(Result::OkUnit(_)) | Some(Result::ErrUnit(_)) => {} + None => bail!("result value at {path} has no case"), + } + } + (SchemaType::Union { spec, .. }, schema_value::Value::UnionValue(union)) => { + let branch = spec + .branches + .iter() + .find(|branch| branch.tag == union.tag) + .ok_or_else(|| anyhow!("unknown union branch '{}' at {path}", union.tag))?; + collect_text_fragments( + graph, + &branch.body, + union + .body + .as_deref() + .ok_or_else(|| anyhow!("union value at {path} has no body"))?, + &format!("{path}.{}", union.tag), + source_language, + fragments, + )?; + } + _ => bail!("invocation value at {path} does not match its declared schema"), + } + Ok(()) +} + +fn discover_streams( + graph: &SchemaGraph, + ty: &SchemaType, + value: &ProtoSchemaValue, + path: &str, + parent_stream_id: Option, + stdout_format: InvocationStdoutFormat, + output: &mut HashMap, +) -> anyhow::Result<()> { + let ty = graph + .resolve_ref(ty) + .map_err(|error| anyhow!(error.to_string()))?; + let Some(value) = value.value.as_ref() else { + bail!("schema value at {path} is empty"); + }; + match (ty, value) { + ( + SchemaType::Stream { + inner: Some(inner), .. + }, + schema_value::Value::StreamReference(reference), + ) => { + let raw_kind = if stdout_format == InvocationStdoutFormat::Raw { + raw_stream_kind(graph, inner) + } else { + None + }; + if output + .insert( + reference.stream_id, + OutputStream { + item_type: (**inner).clone(), + parent_stream_id, + path: path.to_string(), + raw_kind, + next_offset: 0, + terminal: false, + }, + ) + .is_some() + { + bail!( + "output stream {} was discovered more than once", + reference.stream_id + ); + } + } + (SchemaType::Record { fields, .. }, schema_value::Value::RecordValue(record)) => { + for (field, value) in fields.iter().zip(&record.fields) { + discover_streams( + graph, + &field.body, + value, + &format!("{path}.{}", field.name), + parent_stream_id, + stdout_format, + output, + )?; + } + } + (SchemaType::Tuple { elements, .. }, schema_value::Value::TupleValue(tuple)) => { + for (index, (ty, value)) in elements.iter().zip(&tuple.elements).enumerate() { + discover_streams( + graph, + ty, + value, + &format!("{path}[{index}]"), + parent_stream_id, + stdout_format, + output, + )?; + } + } + (SchemaType::List { element, .. }, schema_value::Value::ListValue(list)) => { + for (index, value) in list.elements.iter().enumerate() { + discover_streams( + graph, + element, + value, + &format!("{path}[{index}]"), + parent_stream_id, + stdout_format, + output, + )?; + } + } + (SchemaType::FixedList { element, .. }, schema_value::Value::FixedListValue(list)) => { + for (index, value) in list.elements.iter().enumerate() { + discover_streams( + graph, + element, + value, + &format!("{path}[{index}]"), + parent_stream_id, + stdout_format, + output, + )?; + } + } + ( + SchemaType::Map { + key, + value: value_type, + .. + }, + schema_value::Value::MapValue(map), + ) => { + for (index, entry) in map.entries.iter().enumerate() { + if let Some(value) = &entry.key { + discover_streams( + graph, + key, + value, + &format!("{path}[{index}].key"), + parent_stream_id, + stdout_format, + output, + )?; + } + if let Some(value) = &entry.value { + discover_streams( + graph, + value_type, + value, + &format!("{path}[{index}].value"), + parent_stream_id, + stdout_format, + output, + )?; + } + } + } + (SchemaType::Option { inner, .. }, schema_value::Value::OptionValue(option)) => { + if let Some(value) = &option.inner { + discover_streams( + graph, + inner, + value, + &format!("{path}.some"), + parent_stream_id, + stdout_format, + output, + )?; + } + } + (SchemaType::Variant { cases, .. }, schema_value::Value::VariantValue(variant)) => { + if let Some(case) = cases.get(variant.case as usize) + && let (Some(ty), Some(value)) = (&case.payload, &variant.payload) + { + discover_streams( + graph, + ty, + value, + &format!("{path}.{}", case.name), + parent_stream_id, + stdout_format, + output, + )?; + } + } + (SchemaType::Result { spec, .. }, schema_value::Value::ResultValue(result)) => { + use golem_api_grpc::proto::golem::schema::result_value::Result; + match result.result.as_ref() { + Some(Result::Ok(value)) if spec.ok.is_some() => discover_streams( + graph, + spec.ok.as_deref().unwrap(), + value, + &format!("{path}.ok"), + parent_stream_id, + stdout_format, + output, + )?, + Some(Result::Err(value)) if spec.err.is_some() => discover_streams( + graph, + spec.err.as_deref().unwrap(), + value, + &format!("{path}.err"), + parent_stream_id, + stdout_format, + output, + )?, + _ => {} + } + } + (SchemaType::Union { spec, .. }, schema_value::Value::UnionValue(union)) => { + if let Some(branch) = spec.branches.iter().find(|branch| branch.tag == union.tag) + && let Some(value) = &union.body + { + discover_streams( + graph, + &branch.body, + value, + &format!("{path}.{}", union.tag), + parent_stream_id, + stdout_format, + output, + )?; + } + } + _ if stdout_format == InvocationStdoutFormat::Raw => { + bail!("raw output value at {path} does not match its declared direct stream schema") + } + _ => {} + } + Ok(()) +} + +fn proto_value_to_json( + graph: &SchemaGraph, + ty: &SchemaType, + value: &ProtoSchemaValue, +) -> anyhow::Result { + if let Ok(value) = SchemaValue::try_from(value.clone()) { + return golem_common::schema::render::to_json_value(graph, ty, &value).map_err(Into::into); + } + + let ty = graph + .resolve_ref(ty) + .map_err(|error| anyhow!(error.to_string()))?; + let value = value + .value + .as_ref() + .ok_or_else(|| anyhow!("empty schema value"))?; + Ok(match (ty, value) { + (SchemaType::Stream { .. }, schema_value::Value::StreamReference(reference)) => { + serde_json::json!({ "$stream": reference.stream_id }) + } + (SchemaType::Record { fields, .. }, schema_value::Value::RecordValue(record)) => { + if fields.len() != record.fields.len() { + bail!("record result has the wrong number of fields"); + } + let mut result = serde_json::Map::new(); + for (field, value) in fields.iter().zip(&record.fields) { + result.insert( + field.name.clone(), + proto_value_to_json(graph, &field.body, value)?, + ); + } + serde_json::Value::Object(result) + } + (SchemaType::Tuple { elements, .. }, schema_value::Value::TupleValue(tuple)) => { + if elements.len() != tuple.elements.len() { + bail!("tuple result has the wrong number of elements"); + } + serde_json::Value::Array( + elements + .iter() + .zip(&tuple.elements) + .map(|(ty, value)| proto_value_to_json(graph, ty, value)) + .collect::>()?, + ) + } + (SchemaType::List { element, .. }, schema_value::Value::ListValue(list)) => { + serde_json::Value::Array( + list.elements + .iter() + .map(|value| proto_value_to_json(graph, element, value)) + .collect::>()?, + ) + } + (SchemaType::FixedList { element, .. }, schema_value::Value::FixedListValue(list)) => { + serde_json::Value::Array( + list.elements + .iter() + .map(|value| proto_value_to_json(graph, element, value)) + .collect::>()?, + ) + } + ( + SchemaType::Map { + key, + value: value_type, + .. + }, + schema_value::Value::MapValue(map), + ) => serde_json::Value::Array( + map.entries + .iter() + .map(|entry| { + Ok(serde_json::Value::Array(vec![ + proto_value_to_json( + graph, + key, + entry + .key + .as_ref() + .ok_or_else(|| anyhow!("map entry has no key"))?, + )?, + proto_value_to_json( + graph, + value_type, + entry + .value + .as_ref() + .ok_or_else(|| anyhow!("map entry has no value"))?, + )?, + ])) + }) + .collect::>()?, + ), + (SchemaType::Option { inner, .. }, schema_value::Value::OptionValue(option)) => option + .inner + .as_deref() + .map(|value| proto_value_to_json(graph, inner, value)) + .transpose()? + .unwrap_or(serde_json::Value::Null), + (SchemaType::Variant { cases, .. }, schema_value::Value::VariantValue(variant)) => { + let case = cases + .get(variant.case as usize) + .ok_or_else(|| anyhow!("variant case is out of range"))?; + match (&case.payload, variant.payload.as_deref()) { + (None, None) => case.name.clone().into(), + (Some(ty), Some(value)) => serde_json::json!({ + case.name.clone(): proto_value_to_json(graph, ty, value)? + }), + _ => bail!("variant payload does not match its declared case"), + } + } + (SchemaType::Result { spec, .. }, schema_value::Value::ResultValue(result)) => { + use golem_api_grpc::proto::golem::schema::result_value::Result; + match result.result.as_ref() { + Some(Result::Ok(value)) => serde_json::json!({ + "ok": proto_value_to_json( + graph, + spec.ok.as_deref().ok_or_else(|| anyhow!("unexpected ok payload"))?, + value, + )? + }), + Some(Result::Err(value)) => serde_json::json!({ + "err": proto_value_to_json( + graph, + spec.err.as_deref().ok_or_else(|| anyhow!("unexpected err payload"))?, + value, + )? + }), + Some(Result::OkUnit(_)) => serde_json::json!({ "ok": null }), + Some(Result::ErrUnit(_)) => serde_json::json!({ "err": null }), + None => bail!("result has no case"), + } + } + (SchemaType::Union { spec, .. }, schema_value::Value::UnionValue(union)) => { + let branch = spec + .branches + .iter() + .find(|branch| branch.tag == union.tag) + .ok_or_else(|| anyhow!("unknown union branch '{}'", union.tag))?; + proto_value_to_json( + graph, + &branch.body, + union + .body + .as_deref() + .ok_or_else(|| anyhow!("union has no body"))?, + )? + } + _ => bail!("invocation result does not match its declared schema"), + }) +} + +fn is_connection_closed(error: &tungstenite::Error) -> bool { + matches!( + error, + tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed + ) || matches!(error, tungstenite::Error::Io(error) if error.kind() == ErrorKind::BrokenPipe) +} + +#[cfg(test)] +mod tests { + use super::*; + use golem_common::schema::agent::{InputSchema, NamedField}; + use golem_common::schema::metadata::MetadataEnvelope; + use golem_common::schema::schema_type::{BinaryRestrictions, NamedFieldType}; + use test_r::test; + + fn method(parameters: Vec, output_schema: OutputSchema) -> AgentMethodSchema { + AgentMethodSchema { + name: "test".to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::parameters(parameters), + output_schema, + http_endpoint: Vec::new(), + read_only: None, + } + } + + #[test] + fn raw_stream_kind_resolves_refs() { + let graph = SchemaGraph::empty(); + assert_eq!( + raw_stream_kind(&graph, &SchemaType::u8()), + Some(RawStreamKind::U8) + ); + assert_eq!(raw_stream_kind(&graph, &SchemaType::string()), None); + } + + #[test] + fn value_stdin_parser_preserves_blank_strings() { + let value = parse_method_argument_schema_value( + "", + &SchemaGraph::empty(), + &SchemaType::string(), + &SourceLanguage::Rust, + ) + .unwrap(); + assert_eq!(value, SchemaValue::String(String::new())); + } + + #[test] + fn input_failure_uses_producer_cancellation() { + let request = input_cancel_request( + 1, + 4, + StreamCancelReason::Protocol, + "invalid input".to_string(), + ); + let Some(public_invocation_request::Request::StreamCancel(cancel)) = request.request else { + panic!("expected stream cancellation"); + }; + assert_eq!(cancel.stream_id, 1); + assert_eq!(cancel.offset, 4); + assert_eq!(cancel.role(), StreamCancelRole::InputProducer); + assert_eq!(cancel.reason(), StreamCancelReason::Protocol); + } + + #[test] + fn discovers_nested_output_streams() { + let graph = SchemaGraph::empty(); + let ty = SchemaType::Record { + fields: vec![NamedFieldType { + name: "items".to_string(), + body: SchemaType::stream(Some(SchemaType::u32())), + metadata: MetadataEnvelope::default(), + }], + metadata: MetadataEnvelope::default(), + }; + let value = ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 3 }, + )), + }], + })), + }; + let mut streams = HashMap::new(); + discover_streams( + &graph, + &ty, + &value, + "$", + None, + InvocationStdoutFormat::Value, + &mut streams, + ) + .unwrap(); + assert_eq!(streams.get(&3).unwrap().path, "$.items"); + } + + #[test] + async fn raw_output_rejects_result_that_does_not_match_stream_schema() { + let graph = SchemaGraph::empty(); + let output_schema = + OutputSchema::Single(Box::new(SchemaType::stream(Some(SchemaType::u8())))); + let response = InvocationResponse { + response: Some(invocation_response::Response::Result( + golem_api_grpc::proto::golem::worker::InvocationSessionResult { + result: Some(invocation_session_result::Result::MethodResult( + SchemaValue::U8(1).try_into().unwrap(), + )), + ..Default::default() + }, + )), + }; + let (tx, _rx) = mpsc::channel(1); + let output = OutputChannel { + tx, + interrupt: CancellationToken::new(), + input_failed: CancellationToken::new(), + }; + + let result = handle_response( + response, + &graph, + &output_schema, + &SourceLanguage::Rust, + InvocationStdoutFormat::Raw, + false, + "session-key", + &mut SessionIdentity::default(), + &mut HashMap::new(), + &output, + ) + .await; + + assert!( + result.is_err(), + "raw output must reject a scalar result for a declared stream schema" + ); + } + + #[test] + fn text_rendering_preserves_scalar_fragments_around_streams() { + let graph = SchemaGraph::empty(); + let ty = SchemaType::Tuple { + elements: vec![ + SchemaType::string(), + SchemaType::stream(Some(SchemaType::u32())), + ], + metadata: MetadataEnvelope::default(), + }; + let value = ProtoSchemaValue { + value: Some(schema_value::Value::TupleValue( + golem_api_grpc::proto::golem::schema::TupleValue { + elements: vec![ + SchemaValue::String("metadata".to_string()) + .try_into() + .unwrap(), + ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 3 }, + )), + }, + ], + }, + )), + }; + + assert_eq!( + render_text_fragments(&graph, &ty, &value, "$", &SourceLanguage::Rust).unwrap(), + vec![("$[0]".to_string(), "\"metadata\"".to_string())] + ); + } + + #[test] + fn direct_stream_parameter_binds_stdin_with_odd_client_id() { + let graph = SchemaGraph::empty(); + let method = method( + vec![NamedField::user_supplied( + "values", + SchemaType::stream(Some(SchemaType::u32())), + )], + OutputSchema::Unit, + ); + let (parameters, binding) = prepare_method_parameters( + &graph, + &method, + vec!["-".to_string()], + &SourceLanguage::Rust, + InvocationStdinFormat::Value, + ) + .unwrap(); + + let binding = binding.unwrap(); + assert_eq!(binding.stream_id, 1); + assert_eq!(binding.parameter_name, "values"); + let Some(schema_value::Value::RecordValue(record)) = parameters.value else { + panic!("expected parameter record"); + }; + assert!(matches!( + record.fields[0].value, + Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 1 } + )) + )); + } + + #[test] + fn multiple_and_nested_input_streams_are_rejected() { + let graph = SchemaGraph::empty(); + let direct = method( + vec![ + NamedField::user_supplied("left", SchemaType::stream(Some(SchemaType::u32()))), + NamedField::user_supplied("right", SchemaType::stream(Some(SchemaType::u32()))), + ], + OutputSchema::Unit, + ); + assert!( + prepare_method_parameters( + &graph, + &direct, + vec!["-".to_string(), "-".to_string()], + &SourceLanguage::Rust, + InvocationStdinFormat::Value, + ) + .unwrap_err() + .to_string() + .contains("one stream parameter") + ); + + let nested = method( + vec![NamedField::user_supplied( + "nested", + SchemaType::record(vec![NamedFieldType { + name: "values".to_string(), + body: SchemaType::stream(Some(SchemaType::u32())), + metadata: MetadataEnvelope::default(), + }]), + )], + OutputSchema::Unit, + ); + assert!( + prepare_method_parameters( + &graph, + &nested, + vec!["-".to_string()], + &SourceLanguage::Rust, + InvocationStdinFormat::Value, + ) + .unwrap_err() + .to_string() + .contains("direct stream") + ); + + let nested_item = method( + vec![NamedField::user_supplied( + "nested-item", + SchemaType::stream(Some(SchemaType::record(vec![NamedFieldType { + name: "values".to_string(), + body: SchemaType::stream(Some(SchemaType::u32())), + metadata: MetadataEnvelope::default(), + }]))), + )], + OutputSchema::Unit, + ); + assert!( + prepare_method_parameters( + &graph, + &nested_item, + vec!["-".to_string()], + &SourceLanguage::Rust, + InvocationStdinFormat::Value, + ) + .unwrap_err() + .to_string() + .contains("cannot contain nested streams") + ); + } + + #[test] + fn raw_formats_accept_only_direct_byte_streams() { + let graph = SchemaGraph::empty(); + let invalid_input = method( + vec![NamedField::user_supplied( + "values", + SchemaType::stream(Some(SchemaType::u32())), + )], + OutputSchema::Unit, + ); + assert!( + prepare_method_parameters( + &graph, + &invalid_input, + vec!["-".to_string()], + &SourceLanguage::Rust, + InvocationStdinFormat::Raw, + ) + .unwrap_err() + .to_string() + .contains("stream or stream") + ); + + validate_stdout_format( + &graph, + &OutputSchema::Single(Box::new(SchemaType::stream(Some(SchemaType::binary( + BinaryRestrictions::default(), + ))))), + InvocationStdoutFormat::Raw, + ) + .unwrap(); + assert!( + validate_stdout_format( + &graph, + &OutputSchema::Single(Box::new(SchemaType::record(vec![NamedFieldType { + name: "bytes".to_string(), + body: SchemaType::stream(Some(SchemaType::u8())), + metadata: MetadataEnvelope::default(), + }]))), + InvocationStdoutFormat::Raw, + ) + .unwrap_err() + .to_string() + .contains("one direct stream") + ); + } + + #[test] + fn structured_values_preserve_field_names_and_stream_references() { + let graph = SchemaGraph::empty(); + let ty = SchemaType::record(vec![ + NamedFieldType { + name: "count".to_string(), + body: SchemaType::u32(), + metadata: MetadataEnvelope::default(), + }, + NamedFieldType { + name: "items".to_string(), + body: SchemaType::stream(Some(SchemaType::string())), + metadata: MetadataEnvelope::default(), + }, + ]); + let value = ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ + SchemaValue::U32(2).try_into().unwrap(), + ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 9 }, + )), + }, + ], + })), + }; + + assert_eq!( + proto_value_to_json(&graph, &ty, &value).unwrap(), + serde_json::json!({ "count": 2, "items": { "$stream": 9 } }) + ); + } + + #[test] + fn input_cancellation_is_detected_for_the_bound_stream_only() { + let response = InvocationResponse { + response: Some(invocation_response::Response::StreamCancel( + golem_api_grpc::proto::golem::worker::StreamCancel { + stream_id: 1, + offset: 0, + role: StreamCancelRole::InputConsumer as i32, + reason: StreamCancelReason::Cancelled as i32, + details: None, + }, + )), + }; + assert!(response_cancels_input(&response, Some(1))); + assert!(!response_cancels_input(&response, Some(3))); + } + + #[test] + async fn connection_truncation_before_finish_is_an_error() { + let (wire_tx, _wire_rx) = mpsc::channel(1); + let error = receive_response( + None, + &wire_tx, + &CancellationToken::new(), + &CancellationToken::new(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("ended before completion")); + } + + #[test] + async fn event_after_finish_is_an_error() { + let (wire_tx, _wire_rx) = mpsc::channel(1); + let (_output_tx, mut output_rx) = oneshot::channel(); + let mut frames = futures_util::stream::iter([Ok(Message::Binary(Vec::new().into()))]); + let error = await_clean_close( + &mut frames, + &wire_tx, + &CancellationToken::new(), + &mut output_rx, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("event after completion")); + } + + #[test] + async fn clean_close_wait_is_interrupted() { + let (wire_tx, _wire_rx) = mpsc::channel(1); + let (_output_tx, mut output_rx) = oneshot::channel(); + let mut frames = futures_util::stream::pending(); + let interrupt = CancellationToken::new(); + interrupt.cancel(); + + let error = await_clean_close(&mut frames, &wire_tx, &interrupt, &mut output_rx) + .await + .unwrap_err(); + assert_eq!(error.downcast_ref::().unwrap().0, 130); + } + + #[test] + async fn clean_close_wait_reports_broken_pipe() { + let (wire_tx, _wire_rx) = mpsc::channel(1); + let (output_tx, mut output_rx) = oneshot::channel(); + let mut frames = futures_util::stream::pending(); + output_tx + .send(Err(anyhow!(PipedExitCode(0)))) + .expect("failed to send output error"); + + let error = await_clean_close( + &mut frames, + &wire_tx, + &CancellationToken::new(), + &mut output_rx, + ) + .await + .unwrap_err(); + assert_eq!(error.downcast_ref::().unwrap().0, 0); + } + + #[test] + async fn saturated_wire_send_is_interrupted() { + let (wire_tx, _wire_rx) = mpsc::channel(1); + wire_tx.try_send(Message::Ping(Vec::new().into())).unwrap(); + let interrupt = CancellationToken::new(); + let cancel = interrupt.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + cancel.cancel(); + }); + + let error = tokio::time::timeout( + Duration::from_secs(1), + send_message(&wire_tx, Message::Pong(Vec::new().into()), &interrupt), + ) + .await + .expect("saturated send did not observe cancellation") + .unwrap_err(); + assert_eq!(error.downcast_ref::().unwrap().0, 130); + } + + #[test] + async fn saturated_pong_send_observes_input_failure() { + let (wire_tx, _wire_rx) = mpsc::channel(1); + wire_tx.try_send(Message::Ping(Vec::new().into())).unwrap(); + let input_failed = CancellationToken::new(); + let cancel = input_failed.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + cancel.cancel(); + }); + + let error = tokio::time::timeout( + Duration::from_secs(1), + receive_response( + Some(Ok(Message::Ping(Vec::new().into()))), + &wire_tx, + &CancellationToken::new(), + &input_failed, + ), + ) + .await + .expect("saturated Pong send did not observe input failure") + .unwrap_err(); + assert!(error.downcast_ref::().is_some()); + } + + #[test] + async fn saturated_output_send_observes_input_failure() { + let (tx, _rx) = mpsc::channel(1); + tx.try_send(OutputJob::Text("first".to_string())).unwrap(); + let input_failed = CancellationToken::new(); + let cancel = input_failed.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + cancel.cancel(); + }); + let output = OutputChannel { + tx, + interrupt: CancellationToken::new(), + input_failed, + }; + + let error = tokio::time::timeout( + Duration::from_secs(1), + emit(&output, OutputJob::Text("second".to_string())), + ) + .await + .expect("saturated output send did not observe input failure") + .unwrap_err(); + assert!(error.downcast_ref::().is_some()); + } + + struct BrokenOnFlush { + bytes: Vec, + } + + impl Write for BrokenOnFlush { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + self.bytes.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Err(std::io::Error::new(ErrorKind::BrokenPipe, "closed")) + } + } + + #[test] + fn output_flush_detects_broken_pipe_after_one_raw_item() { + let (output_tx, output_rx) = mpsc::channel(1); + assert!(output_tx.try_send(OutputJob::Raw(vec![1])).is_ok()); + drop(output_tx); + let mut output = BrokenOnFlush { bytes: Vec::new() }; + + let error = write_output_to(output_rx, Format::Text, false, &mut output).unwrap_err(); + assert_eq!(output.bytes, vec![1]); + assert_eq!(error.downcast_ref::().unwrap().0, 0); + } +} diff --git a/cli/golem-cli/src/command_handler/agent/mod.rs b/cli/golem-cli/src/command_handler/agent/mod.rs index ce8bb1adfd..4f8c6816f8 100644 --- a/cli/golem-cli/src/command_handler/agent/mod.rs +++ b/cli/golem-cli/src/command_handler/agent/mod.rs @@ -12,13 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod invocation_session; mod stream; mod stream_output; use crate::command::shared_args::{ AgentFunctionArgument, AgentFunctionName, AgentIdArgs, PostDeployArgs, StreamArgs, }; -use crate::command::worker::AgentSubcommand; +use crate::command::worker::{AgentSubcommand, InvocationStdinFormat, InvocationStdoutFormat}; use crate::command_handler::Handlers; use crate::command_handler::agent::stream::AgentConnection; use crate::context::Context; @@ -26,8 +27,8 @@ use crate::error::NonSuccessfulExit; use crate::error::service::{MapServiceError, ServiceError}; use crate::fuzzy::{Error, FuzzySearch}; use crate::log::{ - LogColorize, LogIndent, log_action, log_error, log_error_action, log_failed_to, log_warn, - log_warn_action, logln, + LogColorize, LogIndent, LogOutput, Output as LogOutputTarget, log_action, log_error, + log_error_action, log_failed_to, log_warn, log_warn_action, logln, }; use crate::model::agent::action_result::{ AgentCancelInvocationResult, AgentDeleteView, AgentFileContentsResult, AgentInterruptResult, @@ -124,6 +125,8 @@ impl AgentCommandHandler { idempotency_key, no_stream, stream_args, + stdin_format, + stdout_format, post_deploy_args, schedule_at, } => { @@ -135,6 +138,8 @@ impl AgentCommandHandler { idempotency_key, no_stream, stream_args, + stdin_format, + stdout_format, post_deploy_args, schedule_at, ) @@ -324,9 +329,13 @@ impl AgentCommandHandler { idempotency_key: Option, no_stream: bool, stream_args: StreamArgs, + stdin_format: InvocationStdinFormat, + stdout_format: InvocationStdoutFormat, post_deploy_args: Option, schedule_at: Option>, ) -> anyhow::Result<()> { + let _raw_output_guard = (stdout_format == InvocationStdoutFormat::Raw) + .then(|| LogOutput::new(LogOutputTarget::None)); self.ctx.silence_app_context_init().await; fn new_idempotency_key() -> IdempotencyKey { @@ -465,6 +474,67 @@ impl AgentCommandHandler { AgentInvocationMode::Await }; + let method_uses_streams = agent_type + .methods + .iter() + .find(|method| method.name == method_name) + .is_some_and(|method| method.uses_streams(&agent_type.schema)); + if !method_uses_streams && stdin_format == InvocationStdinFormat::Raw { + bail!("--stdin-format raw requires a direct stream or stream parameter"); + } + if !method_uses_streams && stdout_format == InvocationStdoutFormat::Raw { + bail!("--stdout-format raw requires a direct stream or stream result"); + } + + if trigger && method_uses_streams { + bail!("Streaming agent methods require an attached invocation session"); + } + + if method_uses_streams { + let mut connect_handle = if !no_stream && stdout_format == InvocationStdoutFormat::Value + { + let connection = AgentConnection::new( + self.ctx.worker_service_url().clone(), + self.ctx.auth_token().await?, + &component.id, + stream_agent_id.to_string(), + stream_args.into(), + self.ctx.allow_insecure(), + self.ctx.format(), + self.ctx.agent_stream_ping_interval(), + Some(idempotency_key.clone()), + ) + .await?; + Some(tokio::spawn(async move { connection.run_forever().await })) + } else { + None + }; + + let result = invocation_session::invoke( + self.ctx.clone(), + invocation_session::InvocationSessionArgs { + application_name: agent_id_match.environment.application_name.to_string(), + environment_name: agent_id_match.environment.environment_name.to_string(), + agent_type, + parsed_agent_id: stream_agent_id, + method_name, + arguments, + config: Vec::new(), + idempotency_key, + stdin_format, + stdout_format, + }, + ) + .await; + + if let Some(mut handle) = connect_handle.take() + && timeout(Duration::from_secs(3), &mut handle).await.is_err() + { + handle.abort(); + } + return result; + } + let source_language = SourceLanguage::from(agent_type.source_language.as_str()); let method_parameters = parse_method_parameters_with_error_table( &agent_type, @@ -2908,7 +2978,7 @@ fn parse_method_parameters_with_error_table( Ok(SchemaValue::Record { fields: values }) } -fn parse_method_argument_schema_value( +pub(super) fn parse_method_argument_schema_value( value: &str, graph: &SchemaGraph, schema: &SchemaType, diff --git a/cli/golem-cli/src/model/agent/invocation_session.rs b/cli/golem-cli/src/model/agent/invocation_session.rs new file mode 100644 index 0000000000..7ce6abc586 --- /dev/null +++ b/cli/golem-cli/src/model/agent/invocation_session.rs @@ -0,0 +1,79 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::model::cli_output::StructuredOutput; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentInvocationSessionEvent { + pub kind: AgentInvocationSessionEventKind, + pub idempotency_key: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub component_revision: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub outcome: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub reason: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_stream_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub path: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub offset: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub value: Option, +} + +impl StructuredOutput for AgentInvocationSessionEvent { + const KIND: &'static str = "agent.invoke-session"; +} + +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum AgentInvocationSessionEventKind { + Accepted, + Rejected, + Result, + Item, + End, + StreamError, + StreamCancel, + Finished, +} + +impl AgentInvocationSessionEvent { + pub fn new(kind: AgentInvocationSessionEventKind, idempotency_key: impl Into) -> Self { + Self { + kind, + idempotency_key: idempotency_key.into(), + agent_id: None, + component_revision: None, + outcome: None, + reason: None, + error: None, + stream_id: None, + parent_stream_id: None, + path: None, + offset: None, + value: None, + } + } +} diff --git a/cli/golem-cli/src/model/agent/mod.rs b/cli/golem-cli/src/model/agent/mod.rs index 4fc36f229d..13667a8f63 100644 --- a/cli/golem-cli/src/model/agent/mod.rs +++ b/cli/golem-cli/src/model/agent/mod.rs @@ -15,6 +15,7 @@ pub mod action_result; pub mod extraction; pub mod files; +pub mod invocation_session; pub mod oplog; pub mod stream; diff --git a/cli/golem-cli/src/model/cli_output/tests.rs b/cli/golem-cli/src/model/cli_output/tests.rs index 31277e5b9f..ef11bb0f73 100644 --- a/cli/golem-cli/src/model/cli_output/tests.rs +++ b/cli/golem-cli/src/model/cli_output/tests.rs @@ -129,6 +129,11 @@ static STRUCTURED_OUTPUT_TEST_REGISTRY: &[StructuredOutputTestEntry] = &[ arb_agent_interrupt_result ), registry_entry!("InvokeResultView", "agent.invoke", arb_agent_invoke_result), + registry_entry!( + "AgentInvocationSessionEvent", + "agent.invoke-session", + arb_agent_invocation_session_event + ), registry_entry!( "AgentsMetadataResponseView", "agent.list", @@ -2821,6 +2826,71 @@ fn arb_agent_stream_event() -> OutputDocumentStrategy { ) } +fn arb_agent_invocation_session_event() -> OutputDocumentStrategy { + serialized_output( + ( + arb_agent_invocation_session_event_kind(), + arb_small_string(), + proptest::option::of(arb_small_string()), + proptest::option::of(arb_small_u64()), + ( + proptest::option::of(arb_small_string()), + proptest::option::of(arb_small_string()), + proptest::option::of(arb_small_string()), + proptest::option::of(arb_small_u64()), + ), + ( + proptest::option::of(arb_small_u64()), + proptest::option::of(arb_small_string()), + proptest::option::of(arb_small_u64()), + proptest::option::of(arb_json_value(1)), + ), + ) + .prop_map( + |( + kind, + idempotency_key, + agent_id, + component_revision, + (outcome, reason, error, stream_id), + (parent_stream_id, path, offset, value), + )| { + crate::model::agent::invocation_session::AgentInvocationSessionEvent { + kind, + idempotency_key, + agent_id, + component_revision, + outcome, + reason, + error, + stream_id, + parent_stream_id, + path, + offset, + value, + } + }, + ), + ) +} + +fn arb_agent_invocation_session_event_kind() +-> BoxedStrategy { + use crate::model::agent::invocation_session::AgentInvocationSessionEventKind; + + prop_oneof![ + Just(AgentInvocationSessionEventKind::Accepted), + Just(AgentInvocationSessionEventKind::Rejected), + Just(AgentInvocationSessionEventKind::Result), + Just(AgentInvocationSessionEventKind::Item), + Just(AgentInvocationSessionEventKind::End), + Just(AgentInvocationSessionEventKind::StreamError), + Just(AgentInvocationSessionEventKind::StreamCancel), + Just(AgentInvocationSessionEventKind::Finished), + ] + .boxed() +} + fn arb_agent_stream_event_kind() -> BoxedStrategy { prop_oneof![ diff --git a/cli/golem-cli/tests/app/agents.rs b/cli/golem-cli/tests/app/agents.rs index 775727e7a2..07a87faf66 100644 --- a/cli/golem-cli/tests/app/agents.rs +++ b/cli/golem-cli/tests/app/agents.rs @@ -1,5 +1,5 @@ use crate::app::{TestContext, cmd, flag, merge_into_manifest}; -use crate::crate_path; +use crate::{crate_path, workspace_path}; use std::path::PathBuf; fn test_data_path() -> PathBuf { @@ -21,6 +21,723 @@ use uuid::Uuid; inherit_test_dep!(Tracing); +async fn streaming_invocation_context() -> TestContext { + let mut ctx = TestContext::new(); + let component_dir = ctx.cwd_path_join("component"); + fs::create_dir_all(component_dir.join("src")).unwrap(); + + let fixture = workspace_path().join("test-components/agent-rpc/golem-it-agent-rpc-rust"); + fs::copy(fixture.join("src/lib.rs"), component_dir.join("src/lib.rs")).unwrap(); + + let sdk_path = workspace_path().join("sdks/rust/golem-rust"); + fs::write_str( + component_dir.join("Cargo.toml"), + formatdoc! {r#" + [package] + name = "golem_it_agent_rpc_rust" + version = "0.0.1" + edition = "2024" + + [profile.release] + opt-level = "s" + lto = true + + [lib] + crate-type = ["cdylib"] + path = "src/lib.rs" + + [dependencies] + bytes = "1.11.0" + log = {{ version = "0.4.29", features = ["kv"] }} + golem-rust = {{ path = "{sdk_path}", features = ["bytes", "export_golem_agentic"] }} + serde = {{ version = "1", features = ["derive"] }} + serde_json = "1" + wasi-fetch = "=0.2.0" + "#, sdk_path = sdk_path.display()}, + ) + .unwrap(); + + fs::write_str( + ctx.cwd_path_join("golem.yaml"), + formatdoc! {r#" + manifestVersion: {MANIFEST_VERSION} + + app: streaming-invocation + + componentTemplates: + rust-streaming-test: + build: + - command: cargo build --target wasm32-wasip2 --release + sources: + - "{{{{ componentDir }}}}/src" + - "{{{{ componentDir }}}}/Cargo.toml" + targets: + - "{{{{ cargoTarget }}}}/wasm32-wasip2/release/golem_it_agent_rpc_rust.wasm" + componentWasm: "{{{{ cargoTarget }}}}/wasm32-wasip2/release/golem_it_agent_rpc_rust.wasm" + outputWasm: "{{{{ golemTempDir }}}}/agents/golem_it_agent_rpc_rust.wasm" + + components: + golem-it:agent-rpc-rust: + dir: component + templates: rust-streaming-test + presets: + release: {{}} + + environments: + local: + server: local + componentPresets: release + "#, MANIFEST_VERSION = versions::sdk::MANIFEST}, + ) + .unwrap(); + + ctx.start_server().await; + let outputs = ctx.cli([cmd::DEPLOY, flag::YES]).await; + assert!(outputs.success_or_dump()); + ctx +} + +fn streaming_agent(name: &str) -> String { + format!(r#"StreamingRpcTarget("{name}")"#) +} + +#[test] +#[timeout("15 minutes")] +async fn test_streaming_invocation_cli_end_to_end() { + let ctx = streaming_invocation_context().await; + let agent = streaming_agent(&Uuid::new_v4().to_string()); + + let scalar = ctx + .cli([cmd::AGENT, cmd::INVOKE, &agent, "increment_scalar"]) + .await; + assert!(scalar.success_or_dump()); + assert!(scalar.stdout_contains("Invocation result in Rust syntax:")); + assert!(scalar.stdout_contains("1")); + + let unit = ctx.cli([cmd::AGENT, cmd::INVOKE, &agent, "noop"]).await; + assert!(unit.success_or_dump()); + assert!(unit.stdout_contains("void")); + + let trigger_stream = ctx + .cli([ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce", + "[1]", + "--trigger", + ]) + .await; + assert!(!trigger_stream.success()); + assert!( + trigger_stream + .stderr_contains("Streaming agent methods require an attached invocation session") + ); + + let scheduled_stream = ctx + .cli([ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce", + "[1]", + "--trigger", + "--schedule-at", + "2030-01-01T00:00:00Z", + ]) + .await; + assert!(!scheduled_stream.success()); + assert!( + scheduled_stream + .stderr_contains("Streaming agent methods require an attached invocation session") + ); + + let value_input = ctx + .cli_with_input( + [cmd::AGENT, cmd::INVOKE, &agent, "consume", "-"], + b"1\n2\n3\n", + ) + .await; + assert!( + value_input.success(), + "value stdin failed: {}", + value_input.stderr_text() + ); + assert!(value_input.stdout_text().contains("[1, 2, 3]")); + + let string_input = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "consume_strings", + "-", + "--no-stream", + ], + b"first\n\nthird\n", + ) + .await; + assert!( + string_input.success(), + "string stdin failed: {}", + string_input.stderr_text() + ); + assert!( + string_input + .stdout_text() + .contains("[\"first\", \"\", \"third\"]\n") + ); + + let value_output = ctx + .cli_with_input( + [cmd::AGENT, cmd::INVOKE, &agent, "produce", "[4, 5, 6]"], + b"", + ) + .await; + assert!( + value_output.success(), + "value stdout failed: {}", + value_output.stderr_text() + ); + assert!(value_output.stdout_text().contains("4\n5\n6\n")); + + let bidirectional = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "transform", + "-", + "--no-stream", + ], + b"7\n8\n", + ) + .await; + assert!( + bidirectional.success(), + "bidirectional invocation failed: {}", + bidirectional.stderr_text() + ); + assert!(bidirectional.stdout_text().contains("70\n80\n")); + + let raw_output = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_bytes", + "[0, 1, 2, 255]", + "--stdout-format", + "raw", + ], + b"", + ) + .await; + assert!( + raw_output.success(), + "raw stdout failed: {}", + raw_output.stderr_text() + ); + assert_eq!(raw_output.stdout(), &[0, 1, 2, 255]); + + let raw_input = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "consume_bytes", + "-", + "--stdin-format", + "raw", + "--no-stream", + ], + b"raw input", + ) + .await; + assert!( + raw_input.success(), + "raw stdin failed: {}", + raw_input.stderr_text() + ); + assert!( + raw_input + .stdout_text() + .contains("[114, 97, 119, 32, 105, 110, 112, 117, 116]") + ); + + let raw_bidirectional = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "transform_bytes", + "-", + "--stdin-format", + "raw", + "--stdout-format", + "raw", + ], + &[0, 9, 10, 255], + ) + .await; + assert!( + raw_bidirectional.success(), + "raw bidirectional invocation failed: {}", + raw_bidirectional.stderr_text() + ); + assert_eq!(raw_bidirectional.stdout(), &[0, 9, 10, 255]); + + let binary_input = (0..(64 * 1024 + 17)) + .map(|index| (index % 251) as u8) + .collect::>(); + let raw_binary = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "transform_binary", + "-", + "--stdin-format", + "raw", + "--stdout-format", + "raw", + ], + &binary_input, + ) + .await; + assert!( + raw_binary.success(), + "raw binary invocation failed: {}", + raw_binary.stderr_text() + ); + assert_eq!(raw_binary.stdout(), binary_input); + + let binary_chunks = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "consume_binary_chunks", + "-", + "--stdin-format", + "raw", + "--no-stream", + ], + &binary_input, + ) + .await; + assert!( + binary_chunks.success(), + "raw binary chunk check failed: {}", + binary_chunks.stderr_text() + ); + assert!(binary_chunks.stdout_text().contains("[65536, 17]\n")); + + let siblings = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_siblings", + "--no-stream", + ], + b"", + ) + .await; + assert!( + siblings.success(), + "sibling output failed: {}", + siblings.stderr_text() + ); + let sibling_stdout = siblings.stdout_text(); + assert!(sibling_stdout.contains("$[0]: \"a\"")); + assert!(sibling_stdout.contains("$[0]: \"b\"")); + assert!(sibling_stdout.contains("$[1]: 63")); + + let nested = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_nested_items", + "--no-stream", + ], + b"", + ) + .await; + assert!( + nested.success(), + "nested output failed: {}", + nested.stderr_text() + ); + let nested_stdout = nested.stdout_text(); + assert!(nested_stdout.contains("first")); + assert!(nested_stdout.contains("second")); + assert!(nested_stdout.contains("$[0].values")); + assert!(nested_stdout.contains("$[1].values")); + assert!(nested_stdout.contains("5")); + + let structured = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce", + "[9, 10]", + flag::FORMAT, + "json", + "--no-stream", + ], + b"", + ) + .await; + assert!( + structured.success(), + "structured invocation failed: {}", + structured.stderr_text() + ); + let events = structured + .stdout_text() + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + assert!( + events + .iter() + .all(|event| event["$type"] == "agent.invoke-session") + ); + let kinds = events + .iter() + .map(|event| event["kind"].as_str().unwrap()) + .collect::>(); + assert_eq!( + kinds, + ["accepted", "result", "item", "item", "end", "finished"] + ); + assert_eq!(events.last().unwrap()["outcome"], "success"); + + let scalar_and_stream = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_scalar_and_stream", + flag::FORMAT, + "json", + "--no-stream", + ], + b"", + ) + .await; + assert!( + scalar_and_stream.success(), + "scalar-plus-stream invocation failed: {}", + scalar_and_stream.stderr_text() + ); + let scalar_and_stream_events = scalar_and_stream + .stdout_text() + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + assert_eq!(scalar_and_stream_events[1]["value"][0], "metadata"); + assert_eq!(scalar_and_stream_events[2]["value"], 11); + assert_eq!(scalar_and_stream_events[3]["value"], 12); + assert_eq!( + scalar_and_stream_events.last().unwrap()["outcome"], + "success" + ); + + let yaml = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce", + "[]", + flag::FORMAT, + "yaml", + "--no-stream", + ], + b"", + ) + .await; + assert!( + yaml.success(), + "YAML invocation failed: {}", + yaml.stderr_text() + ); + assert_eq!(yaml.stdout_text().matches("---\n").count(), 4); + + let toon = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce", + "[]", + flag::FORMAT, + "toon", + "--no-stream", + ], + b"", + ) + .await; + assert!( + toon.success(), + "TOON invocation failed: {}", + toon.stderr_text() + ); + assert_eq!(toon.stdout_text().matches("@toon\n").count(), 4); + assert_eq!(toon.stdout_text().matches("@end\n").count(), 4); + + let malformed_input = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "consume", + "-", + "--no-stream", + ], + b"1\nnot-a-u32\n3\n", + ) + .await; + assert!(!malformed_input.success()); + assert!( + malformed_input + .stderr_text() + .contains("parameter 'input' at line 2") + ); + + let drop_input_agent = streaming_agent(&Uuid::new_v4().to_string()); + let saturated_malformed_input = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &drop_input_agent, + "drop_input", + "-", + "--no-stream", + ], + format!("{}not-a-u32\n", "1\n".repeat(16)).as_bytes(), + ) + .await; + assert!(!saturated_malformed_input.success()); + assert!( + saturated_malformed_input + .stderr_text() + .contains("parameter 'input' at line 17") + ); + + let cancelled_valid_input_agent = streaming_agent(&Uuid::new_v4().to_string()); + let cancelled_valid_input = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &cancelled_valid_input_agent, + "drop_input", + "-", + "--no-stream", + ], + "1\n".repeat(64).as_bytes(), + ) + .await; + assert!( + cancelled_valid_input.success(), + "guest input cancellation failed with queued valid input: {}", + cancelled_valid_input.stderr_text() + ); + assert!(cancelled_valid_input.stdout_text().contains("42")); + + let cancelled_late_malformed_agent = streaming_agent(&Uuid::new_v4().to_string()); + let cancelled_late_malformed_input = format!("{}not-a-u32\n", "1\n".repeat(64)); + let cancelled_late_malformed = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &cancelled_late_malformed_agent, + "drop_input", + "-", + "--no-stream", + ], + cancelled_late_malformed_input.as_bytes(), + ) + .await; + assert!(!cancelled_late_malformed.success()); + assert!( + cancelled_late_malformed + .stderr_text() + .contains("parameter 'input' at line 65") + ); + + let held_input_agent = streaming_agent(&Uuid::new_v4().to_string()); + let held_input = format!("{}not-a-u32\n", "1\n".repeat(32)); + let held_malformed_input = tokio::time::timeout( + Duration::from_secs(10), + ctx.cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &held_input_agent, + "hold_input", + "-", + "--no-stream", + ], + held_input.as_bytes(), + ), + ) + .await + .expect("stdin parse failure remained blocked behind unacknowledged input"); + assert!(!held_malformed_input.success()); + assert!( + held_malformed_input + .stderr_text() + .contains("parameter 'input' at line 33") + ); + + let producer_error = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_error", + "--no-stream", + ], + b"", + ) + .await; + assert!(!producer_error.success()); + assert!(producer_error.stdout_text().lines().any(|line| line == "1")); + assert!(producer_error.stderr_text().contains("Output stream")); + + let sibling_error = ctx + .cli_with_input( + [ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_sibling_error", + "--no-stream", + ], + b"", + ) + .await; + assert!(!sibling_error.success()); + assert!(sibling_error.stderr_text().contains("Output stream")); + assert!(sibling_error.stdout_text().contains("$[1]: 63")); + + let broken_pipe = ctx + .cli_with_broken_stdout([ + cmd::AGENT, + cmd::INVOKE, + &agent, + "produce_byte_then_wait", + "--stdout-format", + "raw", + ]) + .await; + assert!( + broken_pipe.success(), + "broken stdout pipe was treated as failure (exit {:?}): {}", + broken_pipe.exit_code(), + broken_pipe.stderr_text() + ); + assert!(!broken_pipe.stderr_text().contains("Broken pipe")); + + let interrupted_agent = streaming_agent(&Uuid::new_v4().to_string()); + #[cfg(unix)] + { + let exit_code = ctx + .cli_ctrl_c_after( + [ + cmd::AGENT, + cmd::INVOKE, + &interrupted_agent, + "transform", + "-", + flag::FORMAT, + "json", + "--no-stream", + ], + "\"kind\":\"accepted\"", + Duration::ZERO, + ) + .await; + assert_eq!(exit_code, 130); + + let backpressured_agent = streaming_agent(&Uuid::new_v4().to_string()); + let exit_code = ctx + .cli_ctrl_c_after( + [ + cmd::AGENT, + cmd::INVOKE, + &backpressured_agent, + "produce_many_bytes", + "200000", + flag::FORMAT, + "json", + "--no-stream", + ], + "\"kind\":\"accepted\"", + Duration::from_secs(1), + ) + .await; + assert_eq!(exit_code, 130); + } + #[cfg(not(unix))] + { + let interactive_agent = interrupted_agent.clone(); + ctx.cli_interactive( + [ + cmd::AGENT, + cmd::INVOKE, + &interactive_agent, + "transform", + "-", + flag::FORMAT, + "json", + "--no-stream", + ], + |session| { + session.set_expect_timeout(Some(Duration::from_secs(30))); + session.expect_str("\"kind\":\"accepted\"")?; + session.send("\u{3}")?; + session.expect_eof() + }, + ) + .await; + } + + let after_interrupt = ctx + .cli([ + cmd::AGENT, + cmd::INVOKE, + &interrupted_agent, + "increment_scalar", + ]) + .await; + assert!(after_interrupt.success_or_dump()); + assert!(after_interrupt.stdout_contains("1")); +} + #[test] async fn test_rust_counter() { let mut ctx = TestContext::new(); diff --git a/cli/golem-cli/tests/app/mod.rs b/cli/golem-cli/tests/app/mod.rs index ba3138cabb..b33fd0d532 100644 --- a/cli/golem-cli/tests/app/mod.rs +++ b/cli/golem-cli/tests/app/mod.rs @@ -31,6 +31,8 @@ tag_suite!(agents, agents); use crate::{Tracing, crate_path, workspace_path}; use anyhow::Context; use colored::Colorize; +#[cfg(unix)] +use expectrl::process::unix::{Signal, WaitStatus}; use expectrl::{Eof, Expect}; use golem_cli::app::build::task_result_marker::{ ExtractComponentMetadataMarkerHash, TaskResultMarker, @@ -55,7 +57,7 @@ use std::thread::sleep; use std::time::Duration; use tempfile::TempDir; use test_r::{inherit_test_dep, tag_suite}; -use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; use tokio::process::{Child, Command}; use tokio::sync::mpsc; use tokio::time::Instant; @@ -127,6 +129,34 @@ pub struct Output { output: Vec, } +struct RawOutput { + status: ExitStatus, + stdout: Vec, + stderr: Vec, +} + +impl RawOutput { + fn success(&self) -> bool { + self.status.success() + } + + fn exit_code(&self) -> Option { + self.status.code() + } + + fn stdout(&self) -> &[u8] { + &self.stdout + } + + fn stdout_text(&self) -> String { + String::from_utf8_lossy(&self.stdout).into_owned() + } + + fn stderr_text(&self) -> String { + String::from_utf8_lossy(&self.stderr).into_owned() + } +} + impl Output { pub async fn stream_and_collect( quiet: bool, @@ -491,6 +521,121 @@ impl TestContext { .unwrap() } + async fn cli_with_input(&self, command_args: I, input: &[u8]) -> RawOutput + where + I: IntoIterator, + S: AsRef, + { + self.rewrite_local_http_domain_ports(); + + let mut args = vec![ + "--config-dir".to_string(), + self.config_dir.path().to_str().unwrap().to_string(), + ]; + args.extend( + command_args + .into_iter() + .map(|arg| arg.as_ref().to_str().unwrap().to_string()), + ); + + let working_dir = fs::absolute_lexical_path(&self.working_dir).unwrap(); + println!( + "{} {}", + "> working directory:".bold(), + working_dir.display() + ); + println!("{} {}", "> golem-cli".bold(), args.iter().join(" ").blue()); + + let mut child = Command::new(&self.golem_cli_path) + .args(args) + .env_remove("GOLEM_BUILTIN_LOCAL_URL") + .envs(&self.env) + .current_dir(&working_dir) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + + let mut stdin = child.stdin.take().unwrap(); + let mut stdout = child.stdout.take().unwrap(); + let mut stderr = child.stderr.take().unwrap(); + let input = input.to_vec(); + + let write_input = async move { + stdin.write_all(&input).await?; + stdin.shutdown().await + }; + let read_stdout = async move { + let mut bytes = Vec::new(); + stdout.read_to_end(&mut bytes).await?; + Ok::<_, std::io::Error>(bytes) + }; + let read_stderr = async move { + let mut bytes = Vec::new(); + stderr.read_to_end(&mut bytes).await?; + Ok::<_, std::io::Error>(bytes) + }; + + let (input_result, stdout, stderr) = tokio::join!(write_input, read_stdout, read_stderr); + input_result.unwrap(); + + RawOutput { + status: child.wait().await.unwrap(), + stdout: stdout.unwrap(), + stderr: stderr.unwrap(), + } + } + + async fn cli_with_broken_stdout(&self, command_args: I) -> RawOutput + where + I: IntoIterator, + S: AsRef, + { + self.rewrite_local_http_domain_ports(); + + let mut args = vec![ + "--config-dir".to_string(), + self.config_dir.path().to_str().unwrap().to_string(), + ]; + args.extend( + command_args + .into_iter() + .map(|arg| arg.as_ref().to_str().unwrap().to_string()), + ); + + let working_dir = fs::absolute_lexical_path(&self.working_dir).unwrap(); + println!( + "{} {}", + "> working directory:".bold(), + working_dir.display() + ); + println!("{} {}", "> golem-cli".bold(), args.iter().join(" ").blue()); + + let mut child = Command::new(&self.golem_cli_path) + .args(args) + .env_remove("GOLEM_BUILTIN_LOCAL_URL") + .envs(&self.env) + .current_dir(&working_dir) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + + drop(child.stdout.take()); + let mut stderr = child.stderr.take().unwrap(); + + let mut stderr_bytes = Vec::new(); + stderr.read_to_end(&mut stderr_bytes).await.unwrap(); + + RawOutput { + status: child.wait().await.unwrap(), + stdout: Vec::new(), + stderr: stderr_bytes, + } + } + async fn cli_interactive(&self, args: I, session_fn: F) where I: IntoIterator, @@ -551,6 +696,70 @@ impl TestContext { .unwrap() } + #[cfg(unix)] + async fn cli_ctrl_c_after( + &self, + command_args: I, + expected: &str, + pause_after_match: Duration, + ) -> i32 + where + I: IntoIterator, + S: AsRef, + { + self.rewrite_local_http_domain_ports(); + + let mut args = vec![ + "--config-dir".to_string(), + self.config_dir.path().to_str().unwrap().to_string(), + ]; + args.extend( + command_args + .into_iter() + .map(|arg| arg.as_ref().to_str().unwrap().to_string()), + ); + + let working_dir = fs::absolute_lexical_path(&self.working_dir).unwrap(); + let golem_cli_path = self.golem_cli_path.clone(); + let env = self.env.clone(); + let expected = expected.to_string(); + + tokio::task::spawn_blocking(move || { + let mut command = std::process::Command::new(golem_cli_path); + command + .current_dir(working_dir) + .env_remove("GOLEM_BUILTIN_LOCAL_URL") + .envs(env) + .env("TERM", "xterm-256color") + .args(args); + let mut session = expectrl::Session::spawn(command) + .expect("failed to spawn interactive golem-cli session"); + session.set_expect_timeout(Some(Duration::from_secs(30))); + session + .expect(expected.as_str()) + .expect("failed to observe invocation acceptance before Ctrl-C"); + std::thread::sleep(pause_after_match); + session + .get_process_mut() + .signal(Signal::SIGINT) + .expect("failed to send Ctrl-C"); + session + .expect(Eof) + .expect("failed to observe EOF after Ctrl-C"); + + match session + .get_process() + .wait() + .expect("failed to wait for golem-cli exit status") + { + WaitStatus::Exited(_, code) => code, + status => panic!("golem-cli did not exit normally after Ctrl-C: {status:?}"), + } + }) + .await + .expect("failed to run Ctrl-C CLI session") + } + async fn cli_interactive_repl_test(&mut self, args: I, session_fn: F) where I: IntoIterator, diff --git a/cli/golem-cli/tests/bridge_gen/moonbit.rs b/cli/golem-cli/tests/bridge_gen/moonbit.rs index 5d4e76090f..e1b5973b63 100644 --- a/cli/golem-cli/tests/bridge_gen/moonbit.rs +++ b/cli/golem-cli/tests/bridge_gen/moonbit.rs @@ -446,6 +446,7 @@ fn guest_mode_generates_wasm_rpc_project_layout() { assert!(moon_pkg.contains(r#""golemcloud/golem_sdk/agents""#)); assert!(moon_pkg.contains(r#""golemcloud/golem_sdk/rpc""#)); assert!(moon_pkg.contains(r#""golemcloud/golem_sdk/schema_model" @model"#)); + assert!(moon_pkg.contains(r#""golemcloud/golem_sdk/schema_model_host" @model_host"#)); assert!(moon_pkg.contains(r#""golemcloud/golem_sdk/interface/golem/agent/common" @common"#)); assert!(moon_pkg.contains(r#""golemcloud/golem_sdk/interface/golem/core/types" @types"#)); assert!( diff --git a/cli/golem-cli/tests/bridge_gen/scala.rs b/cli/golem-cli/tests/bridge_gen/scala.rs index be7c0f2958..1633fd4a61 100644 --- a/cli/golem-cli/tests/bridge_gen/scala.rs +++ b/cli/golem-cli/tests/bridge_gen/scala.rs @@ -429,7 +429,7 @@ fn guest_wasm_rpc_does_not_emit_external_rest_runtime_references() { /// Guest agent bridges expose the Scala SDK RPC surface: constructors resolve /// remote agents through `RemoteAgentClient`, methods invoke Wasm RPC via -/// `invokeAndAwait`, and the generated Scala.js build is ready for a real +/// `asyncInvokeAndAwait`, and the generated Scala.js build is ready for a real /// compile once the in-tree Scala SDK has been published locally. #[test] fn guest_agent_client_surface_targets_scala_sdk_rpc() { @@ -459,7 +459,7 @@ fn guest_agent_client_surface_targets_scala_sdk_rpc() { assert!(client_source.contains("_root_.scala.Either")); assert!(client_source.contains("CounterAgentRemote")); assert!(client_source.contains("_root_.golem.runtime.rpc.RemoteAgentClient.resolve")); - assert!(client_source.contains("resolved.invokeAndAwait")); + assert!(client_source.contains("resolved.asyncInvokeAndAwait")); assert!(client_source.contains("resolved.cancelableAsyncInvokeAndAwait")); assert!(client_source.contains("_root_.golem.runtime.rpc.CancellationToken")); assert!(client_source.contains("_root_.golem.runtime.rpc.SchemaRpcCodec.encodeValue")); @@ -517,7 +517,7 @@ fn guest_ephemeral_agent_uses_per_invocation_identity() { client_source .contains("_root_.golem.runtime.rpc.CancelableAsyncInvocation[_root_.scala.Unit]") ); - assert!(client_source.contains("resolved.invokeAndAwaitWithMetadata(")); + assert!(client_source.contains("resolved.cancelableAsyncInvokeAndAwaitWithMetadata(")); assert!(client_source.contains("resolved.invokeWithMetadata(")); assert!(client_source.contains("resolved.scheduleInvocationWithMetadata(")); assert!(client_source.contains("resolved.scheduleCancelableInvocationWithMetadata(")); diff --git a/cli/golem-cli/wit/deps/golem-agent/guest.wit b/cli/golem-cli/wit/deps/golem-agent/guest.wit index dbe97ff310..d72413c288 100644 --- a/cli/golem-cli/wit/deps/golem-agent/guest.wit +++ b/cli/golem-cli/wit/deps/golem-agent/guest.wit @@ -15,8 +15,9 @@ interface guest { /// Invokes an agent. If create was not called before, it fails. /// /// `input` is a value tree whose root encodes the method's parameter list. - /// The result is `none` when the method's `output-schema` is `unit`, and - /// `some(value)` for a `single` output. + /// Streams are represented recursively by `stream-value` nodes. The result + /// is `none` when the method's `output-schema` is `unit`, and `some(value)` + /// for a `single` output. invoke: async func(method-name: string, input: schema-value-tree, principal: principal) -> result, agent-error>; /// Gets the agent type. If create was not called before, it fails diff --git a/cli/golem-cli/wit/deps/golem-core-v2/golem-core-v2.wit b/cli/golem-cli/wit/deps/golem-core-v2/golem-core-v2.wit index f393d304d6..2b731300f4 100644 --- a/cli/golem-cli/wit/deps/golem-core-v2/golem-core-v2.wit +++ b/cli/golem-cli/wit/deps/golem-core-v2/golem-core-v2.wit @@ -125,6 +125,17 @@ interface types { /// and reveal it only through capability-gated host interfaces. resource secret; + /// An affine wrapper around a native Component Model stream of schema + /// values. The indirection lets a stream occur anywhere in a recursive + /// `schema-value-tree` while preserving the native stream endpoint. + resource schema-value-stream { + /// Wraps any native schema-value reader for placement in a value tree. + wrap: static async func(reader: stream) -> own; + + /// Consumes a schema-value-stream wrapper and returns its native reader. + unwrap: static async func(value: own) -> stream; + } + // ============================================================ // Schema graph (self-contained type carrier) // ============================================================ @@ -549,8 +560,7 @@ interface types { // Capability nodes secret-value(own), quota-token-handle(own), - - // WASI P3 stubs (parseable only in the schema; no constructible values). + stream-value(own), } record variant-value-payload { @@ -609,4 +619,5 @@ interface types { graph: schema-graph, value: schema-value-tree, } + } diff --git a/cli/golem/src/router.rs b/cli/golem/src/router.rs index 96810028de..1b5c3734d7 100644 --- a/cli/golem/src/router.rs +++ b/cli/golem/src/router.rs @@ -68,6 +68,10 @@ pub async fn start_router( // Worker endpoints .at("/v1/agents/create-agent", worker_service_api.clone()) .at("/v1/agents/invoke-agent", worker_service_api.clone()) + .at( + "/v1/agents/invoke-agent-session", + worker_service_api.clone(), + ) .at( "/v1/components/:component_id/workers", worker_service_api.clone(), diff --git a/docker-examples/published-postgres/nginx.conf.template b/docker-examples/published-postgres/nginx.conf.template index 1aee639480..54d41a6cda 100644 --- a/docker-examples/published-postgres/nginx.conf.template +++ b/docker-examples/published-postgres/nginx.conf.template @@ -39,6 +39,16 @@ server { proxy_pass http://worker-service; } + location = /v1/agents/invoke-agent-session { + proxy_pass http://worker-service; + proxy_http_version 1.1; + proxy_set_header Upgrade "websocket"; + proxy_set_header Connection "upgrade"; + + proxy_read_timeout 3600s; + proxy_send_timeout 3600s; + } + location /v1/agents { proxy_pass http://worker-service; } diff --git a/docs/src/content/next/how-to-guides/moonbit/golem-invoke-agent-moonbit.mdx b/docs/src/content/next/how-to-guides/moonbit/golem-invoke-agent-moonbit.mdx index 2a9909eaca..da7e0581d9 100644 --- a/docs/src/content/next/how-to-guides/moonbit/golem-invoke-agent-moonbit.mdx +++ b/docs/src/content/next/how-to-guides/moonbit/golem-invoke-agent-moonbit.mdx @@ -14,7 +14,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using MoonBit syntax. Multiple return values are rendered as a MoonBit tuple, for example `(1, "ok")`. Methods returning `Unit` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `Unit` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one MoonBit value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consume_values - +cat input.bin | golem agent invoke 'MyAgent()' consume_bytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as MoonBit values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produce_values --stdout-format value +golem agent invoke 'MyAgent()' produce_bytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -73,6 +91,8 @@ golem agent invoke 'staging/MyAgent("user-123")' get_status | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/docs/src/content/next/how-to-guides/rust/golem-invoke-agent-rust.mdx b/docs/src/content/next/how-to-guides/rust/golem-invoke-agent-rust.mdx index d1b01f7e0c..95bc1d0258 100644 --- a/docs/src/content/next/how-to-guides/rust/golem-invoke-agent-rust.mdx +++ b/docs/src/content/next/how-to-guides/rust/golem-invoke-agent-rust.mdx @@ -14,7 +14,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using Rust syntax. Multiple return values are rendered as a Rust tuple, for example `(1, "ok")`. Methods returning `()` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `()` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one Rust value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consume_values - +cat input.bin | golem agent invoke 'MyAgent()' consume_bytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as Rust values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produce_values --stdout-format value +golem agent invoke 'MyAgent()' produce_bytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -68,6 +86,8 @@ golem agent invoke 'staging/MyAgent("user-123")' get_status | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/docs/src/content/next/how-to-guides/scala/golem-invoke-agent-scala.mdx b/docs/src/content/next/how-to-guides/scala/golem-invoke-agent-scala.mdx index efecb3a8da..99f9a7a019 100644 --- a/docs/src/content/next/how-to-guides/scala/golem-invoke-agent-scala.mdx +++ b/docs/src/content/next/how-to-guides/scala/golem-invoke-agent-scala.mdx @@ -14,7 +14,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using Scala syntax. Multiple return values are rendered as a Scala tuple, for example `(1, "ok")`. Methods returning `Unit` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `Unit` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one Scala value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consumeValues - +cat input.bin | golem agent invoke 'MyAgent()' consumeBytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as Scala values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produceValues --stdout-format value +golem agent invoke 'MyAgent()' produceBytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -68,6 +86,8 @@ golem agent invoke 'staging/MyAgent("user-123")' getStatus | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/docs/src/content/next/how-to-guides/ts/golem-invoke-agent-ts.mdx b/docs/src/content/next/how-to-guides/ts/golem-invoke-agent-ts.mdx index afaec4b93d..e80ec44406 100644 --- a/docs/src/content/next/how-to-guides/ts/golem-invoke-agent-ts.mdx +++ b/docs/src/content/next/how-to-guides/ts/golem-invoke-agent-ts.mdx @@ -14,7 +14,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using TypeScript syntax. Multiple return values are rendered as a TypeScript tuple, for example `[1, "ok"]`. Methods returning `void` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `void` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one TypeScript value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf '{ value: 1 }\n{ value: 2 }\n' | golem agent invoke 'MyAgent()' consumeValues - +cat input.bin | golem agent invoke 'MyAgent()' consumeBytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as TypeScript values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produceValues --stdout-format value +golem agent invoke 'MyAgent()' produceBytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -68,6 +86,8 @@ golem agent invoke 'staging/MyAgent("user-123")' getStatus | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/docs/src/content/next/rest-api/agent.mdx b/docs/src/content/next/rest-api/agent.mdx index e22a7f87dc..e62ff6bd2f 100644 --- a/docs/src/content/next/rest-api/agent.mdx +++ b/docs/src/content/next/rest-api/agent.mdx @@ -99,6 +99,19 @@ Path|Method|Protected } ``` +## Invoke an agent through an attached live streaming session +Path|Method|Protected +---|---|--- +`/v1/agents/invoke-agent-session`|GET|Yes + + + + + + + + + ## undefined Path|Method|Protected ---|---|--- diff --git a/golem-api-grpc/proto/golem/schema/schema.proto b/golem-api-grpc/proto/golem/schema/schema.proto index f7bfe3b912..4a25c64c52 100644 --- a/golem-api-grpc/proto/golem/schema/schema.proto +++ b/golem-api-grpc/proto/golem/schema/schema.proto @@ -330,9 +330,17 @@ message SchemaValue { // Capability nodes SecretValue secret_value = 70; QuotaTokenValue quota_token_value = 71; + + // Live invocation only. The ID is scoped to one live value session and + // must be resolved by that session rather than by generic protobuf decoding. + SchemaValueStreamReference stream_reference = 80; } } +message SchemaValueStreamReference { + uint64 stream_id = 1; +} + message RecordValue { repeated SchemaValue fields = 1; } diff --git a/golem-api-grpc/proto/golem/worker/invocation_session.proto b/golem-api-grpc/proto/golem/worker/invocation_session.proto new file mode 100644 index 0000000000..88c128f3b5 --- /dev/null +++ b/golem-api-grpc/proto/golem/worker/invocation_session.proto @@ -0,0 +1,214 @@ +syntax = "proto3"; + +package golem.worker; + +import "golem/auth/auth_ctx.proto"; +import "golem/common/account_id.proto"; +import "golem/common/empty.proto"; +import "golem/common/environment.proto"; +import "golem/common/uuid.proto"; +import "golem/component/agent.proto"; +import "golem/schema/schema.proto"; +import "golem/worker/agent_config.proto"; +import "golem/worker/idempotency_key.proto"; +import "golem/worker/invocation.proto"; +import "golem/worker/invocation_context.proto"; +import "golem/worker/v1/worker_execution_error.proto"; +import "golem/worker/worker_id.proto"; +import "google/protobuf/timestamp.proto"; + +message InvocationStart { + golem.worker.AgentId agent_id = 1; + optional string method_name = 2; + golem.schema.SchemaValue input = 3; + golem.worker.IdempotencyKey idempotency_key = 4; + golem.worker.InvocationContext context = 5; + golem.auth.AuthCtx auth_ctx = 6; + golem.component.Principal principal = 7; + golem.common.EnvironmentId environment_id = 8; + repeated golem.worker.AgentConfigEntryDto config = 9; + golem.common.AccountId component_owner_account_id = 10; + golem.worker.AgentInvocationMode mode = 11; + optional google.protobuf.Timestamp schedule_at = 12; + InvocationFreshnessDisposition freshness_disposition = 13; +} + +message PublicInvocationStart { + string application_name = 1; + string environment_name = 2; + string agent_type_name = 3; + golem.schema.SchemaValue constructor_parameters = 4; + optional golem.common.UUID phantom_id = 5; + repeated golem.worker.AgentConfigEntryDto config = 6; + string method_name = 7; + golem.schema.SchemaValue method_parameters = 8; + golem.worker.IdempotencyKey idempotency_key = 9; +} + +message ResumeAttach { + golem.worker.IdempotencyKey idempotency_key = 1; +} + +enum InvocationFreshnessDisposition { + INVOCATION_FRESHNESS_DISPOSITION_MAY_EXIST = 0; + INVOCATION_FRESHNESS_DISPOSITION_KNOWN_FRESH = 1; +} + +message InputStreamItem { + uint64 stream_id = 1; + uint64 sequence = 2; + oneof payload { + golem.schema.SchemaValue value = 3; + bytes packed_u8 = 4; + } +} + +message InputStreamEnd { + uint64 stream_id = 1; + uint64 offset = 2; +} + +message OutputStreamItem { + uint64 stream_id = 1; + uint64 offset = 2; + golem.schema.SchemaValue value = 3; +} + +message OutputStreamEnd { + uint64 stream_id = 1; + uint64 offset = 2; +} + +message OutputStreamError { + uint64 stream_id = 1; + uint64 offset = 2; + string details = 3; +} + +enum StreamCancelRole { + STREAM_CANCEL_ROLE_UNSPECIFIED = 0; + STREAM_CANCEL_ROLE_INPUT_PRODUCER = 1; + STREAM_CANCEL_ROLE_INPUT_CONSUMER = 2; + STREAM_CANCEL_ROLE_OUTPUT_PRODUCER = 3; + STREAM_CANCEL_ROLE_OUTPUT_CONSUMER = 4; +} + +enum StreamCancelReason { + STREAM_CANCEL_REASON_UNSPECIFIED = 0; + STREAM_CANCEL_REASON_CANCELLED = 1; + STREAM_CANCEL_REASON_TRANSPORT = 2; + STREAM_CANCEL_REASON_PROTOCOL = 3; +} + +message StreamCancel { + uint64 stream_id = 1; + uint64 offset = 2; + StreamCancelRole role = 3; + StreamCancelReason reason = 4; + optional string details = 5; +} + +message InputStreamAck { + uint64 stream_id = 1; + uint64 sequence = 2; + uint64 logical_item_count = 3; +} + +message InvocationAccepted { + golem.worker.AgentId agent_id = 1; + golem.worker.IdempotencyKey idempotency_key = 2; + optional uint64 component_revision = 3; +} + +message InvocationSessionResult { + oneof result { + golem.schema.SchemaValue method_result = 1; + golem.common.Empty no_result = 2; + } + optional uint64 component_revision = 3; + golem.worker.AgentId agent_id = 4; + golem.worker.IdempotencyKey idempotency_key = 5; + optional uint64 fuel_consumed = 6; + optional golem.worker.InvocationStatus status = 7; + optional uint64 oplog_index = 8; + optional golem.common.UUID agent_fingerprint = 9; +} + +enum InvocationRejectionReason { + INVOCATION_REJECTION_REASON_UNSPECIFIED = 0; + INVOCATION_REJECTION_REASON_VALIDATION = 1; + INVOCATION_REJECTION_REASON_UNAUTHORIZED = 2; + INVOCATION_REJECTION_REASON_NOT_FOUND = 3; + INVOCATION_REJECTION_REASON_RESUME_UNSUPPORTED = 4; + INVOCATION_REJECTION_REASON_PROTOCOL = 5; + INVOCATION_REJECTION_REASON_INTERNAL = 6; +} + +message InvocationRejected { + InvocationRejectionReason reason = 1; + string error = 2; + golem.worker.IdempotencyKey idempotency_key = 3; + golem.worker.AgentId agent_id = 4; + optional uint64 component_revision = 5; +} + +message AttachmentRevoked { + string details = 1; +} + +enum InvocationFailureKind { + INVOCATION_FAILURE_KIND_UNSPECIFIED = 0; + INVOCATION_FAILURE_KIND_EXECUTION = 1; + INVOCATION_FAILURE_KIND_PROTOCOL = 2; + INVOCATION_FAILURE_KIND_INTERNAL = 3; + INVOCATION_FAILURE_KIND_TRANSPORT = 4; +} + +message InvocationFailure { + InvocationFailureKind kind = 1; + string code = 2; + string message = 3; + optional golem.worker.v1.WorkerExecutionError worker_error = 4; +} + +message InvocationSessionCompletion { + oneof outcome { + golem.common.Empty success = 1; + InvocationFailure failure = 2; + } +} + +message PublicInvocationRequest { + oneof request { + PublicInvocationStart start = 1; + ResumeAttach resume_attach = 2; + InputStreamItem input_item = 3; + InputStreamEnd input_end = 4; + StreamCancel stream_cancel = 5; + } +} + +message InvocationRequest { + oneof request { + InvocationStart start = 1; + ResumeAttach resume_attach = 2; + InputStreamItem input_item = 3; + InputStreamEnd input_end = 4; + StreamCancel stream_cancel = 5; + } +} + +message InvocationResponse { + oneof response { + InvocationAccepted accepted = 1; + InvocationRejected rejected = 2; + InvocationSessionResult result = 3; + OutputStreamItem output_item = 4; + OutputStreamEnd output_end = 5; + OutputStreamError output_error = 6; + InputStreamAck input_ack = 7; + StreamCancel stream_cancel = 8; + AttachmentRevoked attachment_revoked = 9; + InvocationSessionCompletion finished = 10; + } +} diff --git a/golem-api-grpc/proto/golem/worker/v1/worker_service.proto b/golem-api-grpc/proto/golem/worker/v1/worker_service.proto index ff343d1c2d..3d11f922bc 100644 --- a/golem-api-grpc/proto/golem/worker/v1/worker_service.proto +++ b/golem-api-grpc/proto/golem/worker/v1/worker_service.proto @@ -18,6 +18,7 @@ import "golem/worker/filesystem.proto"; import "golem/worker/idempotency_key.proto"; import "golem/worker/invocation.proto"; import "golem/worker/invocation_context.proto"; +import "golem/worker/invocation_session.proto"; import "golem/worker/log_event.proto"; import "golem/worker/oplog_cursor.proto"; import "golem/worker/public_oplog.proto"; @@ -37,6 +38,9 @@ service WorkerService { rpc RevertWorker(RevertWorkerRequest) returns (RevertWorkerResponse); rpc CompletePromise (CompletePromiseRequest) returns (CompletePromiseResponse); rpc InvokeAgent (InvokeAgentRequest) returns (InvokeAgentResponse); + // Internal executor-routing endpoint. Public clients use the authenticated + // agent-session WebSocket and its PublicInvocationRequest envelope. + rpc InvokeAgentSession (stream golem.worker.InvocationRequest) returns (stream golem.worker.InvocationResponse); rpc CancelInvocation (CancelInvocationRequest) returns (CancelInvocationResponse); rpc ProcessOplogEntries (ProcessOplogEntriesRequest) returns (ProcessOplogEntriesResponse); } diff --git a/golem-api-grpc/proto/golem/workerexecutor/v1/worker_executor.proto b/golem-api-grpc/proto/golem/workerexecutor/v1/worker_executor.proto index d97b15fc69..624e6863d0 100644 --- a/golem-api-grpc/proto/golem/workerexecutor/v1/worker_executor.proto +++ b/golem-api-grpc/proto/golem/workerexecutor/v1/worker_executor.proto @@ -10,14 +10,13 @@ import "golem/common/revert_worker_response.proto"; import "golem/common/uuid.proto"; import "golem/component/agent.proto"; import "golem/component/component_id.proto"; -import "golem/schema/schema.proto"; import "golem/shardmanager/shard_id.proto"; import "golem/worker/agent_config.proto"; import "golem/worker/cursor.proto"; import "golem/worker/filesystem.proto"; import "golem/worker/idempotency_key.proto"; -import "golem/worker/invocation.proto"; import "golem/worker/invocation_context.proto"; +import "golem/worker/invocation_session.proto"; import "golem/worker/log_event.proto"; import "golem/worker/oplog_cursor.proto"; import "golem/worker/promise_id.proto"; @@ -29,7 +28,6 @@ import "golem/worker/worker_filter.proto"; import "golem/worker/worker_id.proto"; import "golem/worker/worker_metadata.proto"; import "golem/worker/worker_status.proto"; -import "google/protobuf/timestamp.proto"; service WorkerExecutor { rpc CreateWorker(CreateWorkerRequest) returns (CreateWorkerResponse); @@ -59,7 +57,7 @@ service WorkerExecutor { rpc ActivatePlugin(ActivatePluginRequest) returns (ActivatePluginResponse); rpc DeactivatePlugin(DeactivatePluginRequest) returns (DeactivatePluginResponse); - rpc InvokeAgent(InvokeAgentRequest) returns (InvokeAgentResponse); + rpc InvokeAgentSession(stream golem.worker.InvocationRequest) returns (stream golem.worker.InvocationResponse); rpc ProcessOplogEntries(ProcessOplogEntriesRequest) returns (ProcessOplogEntriesResponse); } @@ -466,52 +464,6 @@ message CancelInvocationResponse { } } -message InvokeAgentRequest { - golem.worker.AgentId agent_id = 1; - optional string method_name = 5; - optional golem.schema.SchemaValue method_parameters = 6; - golem.worker.AgentInvocationMode mode = 7; - optional google.protobuf.Timestamp schedule_at = 8; - optional golem.worker.IdempotencyKey idempotency_key = 9; - golem.common.AccountId component_owner_account_id = 10; - golem.common.EnvironmentId environment_id = 11; - golem.auth.AuthCtx auth_ctx = 12; - optional golem.worker.InvocationContext context = 13; - golem.component.Principal principal = 14; - InvocationFreshnessDisposition freshness_disposition = 15; - repeated golem.worker.AgentConfigEntryDto config = 16; -} - -enum InvocationFreshnessDisposition { - INVOCATION_FRESHNESS_DISPOSITION_MAY_EXIST = 0; - INVOCATION_FRESHNESS_DISPOSITION_KNOWN_FRESH = 1; -} - -message InvokeAgentResponse { - oneof result { - InvokeAgentSuccess success = 1; - golem.worker.v1.WorkerExecutionError failure = 2; - } -} - -message InvokeAgentSuccess { - optional golem.schema.SchemaValue result = 1; - optional uint64 fuel_consumed = 2; - optional uint64 component_revision = 3; - optional golem.worker.InvocationStatus status = 4; - // Oplog index of the agent right after this invocation completed. - optional uint64 oplog_index = 5; - // Per-instance agent fingerprint (UUID) of the agent that produced this - // invocation result. - optional golem.common.UUID agent_fingerprint = 6; - // Final invocation target. For ephemeral agents this contains the generated - // one-shot phantom ID and can be used by observation and control APIs. - optional golem.worker.AgentId agent_id = 7; - // Final invocation key. Together with agent_id this identifies the logical - // invocation across retries and asynchronous result lookup. - optional golem.worker.IdempotencyKey idempotency_key = 8; -} - message ProcessOplogEntriesRequest { golem.worker.AgentId agent_id = 1; golem.common.EnvironmentId environment_id = 2; diff --git a/golem-api-grpc/src/invocation_session_protocol.rs b/golem-api-grpc/src/invocation_session_protocol.rs new file mode 100644 index 0000000000..07b06aafda --- /dev/null +++ b/golem-api-grpc/src/invocation_session_protocol.rs @@ -0,0 +1,1596 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::proto::golem::schema::{SchemaValue, result_value, schema_value}; +use crate::proto::golem::worker::input_stream_item::Payload; +use crate::proto::golem::worker::{ + AgentId, IdempotencyKey, InputStreamAck, InputStreamEnd, InputStreamItem, InvocationAccepted, + InvocationFailureKind, InvocationRejectionReason, InvocationRequest, InvocationResponse, + InvocationSessionCompletion, InvocationSessionResult, PublicInvocationRequest, StreamCancel, + StreamCancelReason, StreamCancelRole, invocation_request, invocation_response, + invocation_session_completion, invocation_session_result, public_invocation_request, +}; +use std::collections::{HashMap, HashSet, VecDeque}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SessionPhase { + Initial, + AwaitDecision { resume: bool }, + Active, + Complete, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct PendingAck { + sequence: u64, + logical_item_count: u64, +} + +#[derive(Debug, Default)] +struct InputState { + next_offset: u64, + terminal: bool, + discard_next_offset: Option, + pending_acks: VecDeque, +} + +#[derive(Debug, Default)] +struct OutputState { + next_offset: u64, + terminal: bool, + cancellation_requested: Option, +} + +#[derive(Debug)] +pub struct InvocationSessionState { + phase: SessionPhase, + idempotency_key: Option, + accepted_agent_id: Option, + accepted_revision: Option, + has_result: bool, + inputs: HashMap, + outputs: HashMap, +} + +impl Default for InvocationSessionState { + fn default() -> Self { + Self { + phase: SessionPhase::Initial, + idempotency_key: None, + accepted_agent_id: None, + accepted_revision: None, + has_result: false, + inputs: HashMap::new(), + outputs: HashMap::new(), + } + } +} + +enum RequestMessage<'a> { + Start { + idempotency_key: &'a Option, + input: Option<&'a SchemaValue>, + }, + ResumeAttach { + idempotency_key: &'a Option, + }, + InputItem(&'a InputStreamItem), + InputEnd(&'a InputStreamEnd), + StreamCancel(&'a StreamCancel), +} + +impl InvocationSessionState { + pub fn validate_public_request( + &mut self, + request: &PublicInvocationRequest, + ) -> Result<(), String> { + self.validate_public_request_with_terminal_race(request, false) + } + + /// Validates a public request at the receiving end of a full-duplex session. + /// + /// A consumer may send an output cancellation before observing a terminal response that the + /// receiver has already recorded. This accepts that cancellation once while preserving strict + /// validation for locally generated requests. + pub fn validate_received_public_request( + &mut self, + request: &PublicInvocationRequest, + ) -> Result<(), String> { + self.validate_public_request_with_terminal_race(request, true) + } + + fn validate_public_request_with_terminal_race( + &mut self, + request: &PublicInvocationRequest, + accept_terminal_output_cancellation: bool, + ) -> Result<(), String> { + let message = + match request.request.as_ref() { + Some(public_invocation_request::Request::Start(start)) => { + if start.application_name.is_empty() + || start.environment_name.is_empty() + || start.agent_type_name.is_empty() + || start.method_name.is_empty() + { + return Err("public invocation selectors must not be empty".to_string()); + } + if start.constructor_parameters.is_none() { + return Err("public invocation has no constructor parameters".to_string()); + } + RequestMessage::Start { + idempotency_key: &start.idempotency_key, + input: Some(start.method_parameters.as_ref().ok_or_else(|| { + "public invocation has no method parameters".to_string() + })?), + } + } + Some(public_invocation_request::Request::ResumeAttach(resume)) => { + RequestMessage::ResumeAttach { + idempotency_key: &resume.idempotency_key, + } + } + Some(public_invocation_request::Request::InputItem(item)) => { + RequestMessage::InputItem(item) + } + Some(public_invocation_request::Request::InputEnd(end)) => { + RequestMessage::InputEnd(end) + } + Some(public_invocation_request::Request::StreamCancel(cancel)) => { + RequestMessage::StreamCancel(cancel) + } + None => return Err("invocation request has no payload".to_string()), + }; + self.validate_request(message, accept_terminal_output_cancellation) + } + + pub fn validate_trusted_request(&mut self, request: &InvocationRequest) -> Result<(), String> { + self.validate_trusted_request_with_terminal_race(request, false) + } + + /// Validates a trusted request at the receiving end of a full-duplex session. + /// + /// A consumer may send an output cancellation before observing a terminal response that the + /// receiver has already recorded. This accepts that cancellation once while preserving strict + /// validation for locally generated requests. + pub fn validate_received_trusted_request( + &mut self, + request: &InvocationRequest, + ) -> Result<(), String> { + self.validate_trusted_request_with_terminal_race(request, true) + } + + fn validate_trusted_request_with_terminal_race( + &mut self, + request: &InvocationRequest, + accept_terminal_output_cancellation: bool, + ) -> Result<(), String> { + let message = match request.request.as_ref() { + Some(invocation_request::Request::Start(start)) => RequestMessage::Start { + idempotency_key: &start.idempotency_key, + input: start.input.as_ref(), + }, + Some(invocation_request::Request::ResumeAttach(resume)) => { + RequestMessage::ResumeAttach { + idempotency_key: &resume.idempotency_key, + } + } + Some(invocation_request::Request::InputItem(item)) => RequestMessage::InputItem(item), + Some(invocation_request::Request::InputEnd(end)) => RequestMessage::InputEnd(end), + Some(invocation_request::Request::StreamCancel(cancel)) => { + RequestMessage::StreamCancel(cancel) + } + None => return Err("invocation request has no payload".to_string()), + }; + self.validate_request(message, accept_terminal_output_cancellation) + } + + pub fn validate_response(&mut self, response: &InvocationResponse) -> Result<(), String> { + if self.phase == SessionPhase::Complete { + return Err("invocation response received after completion".to_string()); + } + let response = response + .response + .as_ref() + .ok_or_else(|| "invocation response has no payload".to_string())?; + + match (self.phase, response) { + ( + SessionPhase::AwaitDecision { resume: false }, + invocation_response::Response::Accepted(accepted), + ) => self.accept(accepted), + ( + SessionPhase::AwaitDecision { resume }, + invocation_response::Response::Rejected(rejected), + ) => { + let reason = + InvocationRejectionReason::try_from(rejected.reason).map_err(|_| { + format!("invalid invocation rejection reason {}", rejected.reason) + })?; + if reason == InvocationRejectionReason::Unspecified { + return Err("invocation rejection reason is unspecified".to_string()); + } + if resume && reason != InvocationRejectionReason::ResumeUnsupported { + return Err("resume-attach must be rejected as resume-unsupported".to_string()); + } + self.validate_idempotency_key(&rejected.idempotency_key)?; + self.phase = SessionPhase::Complete; + Ok(()) + } + (SessionPhase::AwaitDecision { resume: true }, _) => { + Err("resume-attach must receive invocation-rejected".to_string()) + } + (SessionPhase::AwaitDecision { resume: false }, _) => { + Err("invocation must be accepted or rejected before other responses".to_string()) + } + (SessionPhase::Active, invocation_response::Response::Accepted(_)) => { + Err("invocation response contains more than one acceptance".to_string()) + } + (SessionPhase::Active, invocation_response::Response::Rejected(_)) => { + Err("invocation rejection may only appear before acceptance".to_string()) + } + (SessionPhase::Active, invocation_response::Response::Result(result)) => { + self.validate_result(result) + } + (SessionPhase::Active, invocation_response::Response::OutputItem(item)) => { + let value = item + .value + .as_ref() + .ok_or_else(|| "output stream item has no value".to_string())?; + let state = self + .outputs + .get(&item.stream_id) + .ok_or_else(|| format!("output stream {} is unknown", item.stream_id))?; + ensure_open(state.terminal, item.stream_id)?; + if item.offset != state.next_offset { + return Err(format!( + "output stream {} expected offset {}, got {}", + item.stream_id, state.next_offset, item.offset + )); + } + let discovered = stream_references(value)?; + self.ensure_new_output_streams(&discovered)?; + let state = self.outputs.get_mut(&item.stream_id).unwrap(); + state.next_offset = state + .next_offset + .checked_add(1) + .ok_or_else(|| format!("output stream {} offset overflow", item.stream_id))?; + self.insert_output_streams(discovered); + Ok(()) + } + (SessionPhase::Active, invocation_response::Response::OutputEnd(end)) => { + self.terminate_output(end.stream_id, end.offset) + } + (SessionPhase::Active, invocation_response::Response::OutputError(error)) => { + self.terminate_output(error.stream_id, error.offset) + } + (SessionPhase::Active, invocation_response::Response::InputAck(ack)) => { + self.validate_ack(ack) + } + (SessionPhase::Active, invocation_response::Response::StreamCancel(cancel)) => { + validate_cancel(cancel)?; + match cancel.role() { + StreamCancelRole::InputConsumer => { + self.cancel_input(cancel.stream_id, cancel.offset, true) + } + StreamCancelRole::OutputProducer => { + self.confirm_output_cancellation(cancel.stream_id, cancel.offset) + } + _ => Err( + "server response may only cancel an input consumer or output producer" + .to_string(), + ), + } + } + (SessionPhase::Active, invocation_response::Response::AttachmentRevoked(_)) => { + Err("attachment-revoked is not supported by GOL-91".to_string()) + } + (SessionPhase::Active, invocation_response::Response::Finished(finished)) => { + self.finish(finished) + } + (SessionPhase::Initial, _) => { + Err("invocation response received before the first request".to_string()) + } + (SessionPhase::Complete, _) => unreachable!(), + } + } + + pub fn is_complete(&self) -> bool { + self.phase == SessionPhase::Complete + } + + fn validate_request( + &mut self, + message: RequestMessage<'_>, + accept_terminal_output_cancellation: bool, + ) -> Result<(), String> { + match (self.phase, message) { + ( + SessionPhase::Initial, + RequestMessage::Start { + idempotency_key, + input, + }, + ) => self.start(idempotency_key, input, false), + (SessionPhase::Initial, RequestMessage::ResumeAttach { idempotency_key }) => { + self.start(idempotency_key, None, true) + } + (SessionPhase::Initial, _) => { + Err("the first invocation request must be start or resume-attach".to_string()) + } + (SessionPhase::AwaitDecision { .. }, _) => { + Err("invocation input may only be sent after acceptance".to_string()) + } + ( + SessionPhase::Active, + RequestMessage::Start { .. } | RequestMessage::ResumeAttach { .. }, + ) => Err( + "invocation start or resume-attach may only appear as the first request" + .to_string(), + ), + (SessionPhase::Active, RequestMessage::InputItem(item)) => { + self.validate_input_item(item) + } + (SessionPhase::Active, RequestMessage::InputEnd(end)) => { + self.terminate_input(end.stream_id, end.offset) + } + (SessionPhase::Active, RequestMessage::StreamCancel(cancel)) => { + validate_cancel(cancel)?; + match cancel.role() { + StreamCancelRole::InputProducer => { + self.cancel_input(cancel.stream_id, cancel.offset, false) + } + StreamCancelRole::OutputConsumer => self.request_output_cancellation( + cancel.stream_id, + cancel.offset, + accept_terminal_output_cancellation, + ), + _ => Err( + "client request may only cancel an input producer or output consumer" + .to_string(), + ), + } + } + (SessionPhase::Complete, _) => { + Err("invocation request received after completion".to_string()) + } + } + } + + fn start( + &mut self, + idempotency_key: &Option, + input: Option<&SchemaValue>, + resume: bool, + ) -> Result<(), String> { + let key = required_idempotency_key(idempotency_key)?; + let stream_ids = input + .map(stream_references) + .transpose()? + .unwrap_or_default(); + self.inputs = stream_ids + .into_iter() + .map(|stream_id| (stream_id, InputState::default())) + .collect(); + self.idempotency_key = Some(key.to_string()); + self.phase = SessionPhase::AwaitDecision { resume }; + Ok(()) + } + + fn accept(&mut self, accepted: &InvocationAccepted) -> Result<(), String> { + self.validate_idempotency_key(&accepted.idempotency_key)?; + let agent_id = accepted + .agent_id + .as_ref() + .ok_or_else(|| "invocation acceptance has no agent identity".to_string())?; + self.accepted_agent_id = Some(agent_id.clone()); + self.accepted_revision = accepted.component_revision; + self.phase = SessionPhase::Active; + Ok(()) + } + + fn validate_result(&mut self, result: &InvocationSessionResult) -> Result<(), String> { + if self.has_result { + return Err("invocation response contains more than one result".to_string()); + } + self.validate_identity(&result.agent_id, &result.idempotency_key)?; + if let (Some(accepted), Some(result)) = (self.accepted_revision, result.component_revision) + && accepted != result + { + return Err(format!( + "invocation result revision {result} differs from accepted revision {accepted}" + )); + } + let discovered = match result.result.as_ref() { + Some(invocation_session_result::Result::MethodResult(value)) => { + stream_references(value)? + } + Some(invocation_session_result::Result::NoResult(_)) => Vec::new(), + None => return Err("invocation result has no value".to_string()), + }; + self.ensure_new_output_streams(&discovered)?; + self.insert_output_streams(discovered); + self.has_result = true; + Ok(()) + } + + fn validate_input_item(&mut self, item: &InputStreamItem) -> Result<(), String> { + let (logical_item_count, discovered) = match item.payload.as_ref() { + Some(Payload::Value(value)) => (1, stream_references(value)?), + Some(Payload::PackedU8(bytes)) if !bytes.is_empty() => (bytes.len() as u64, Vec::new()), + Some(Payload::PackedU8(_)) => { + return Err("packed-u8 input item must not be empty".to_string()); + } + None => return Err("input stream item has no payload".to_string()), + }; + self.ensure_new_input_streams(&discovered)?; + let state = self + .inputs + .get_mut(&item.stream_id) + .ok_or_else(|| format!("input stream {} is unknown", item.stream_id))?; + if let Some(discard_next_offset) = state.discard_next_offset.as_mut() { + if item.sequence != *discard_next_offset { + return Err(format!( + "cancelled input stream {} expected discarded sequence {}, got {}", + item.stream_id, *discard_next_offset, item.sequence + )); + } + *discard_next_offset = discard_next_offset + .checked_add(logical_item_count) + .ok_or_else(|| format!("input stream {} offset overflow", item.stream_id))?; + self.insert_input_streams(discovered); + return Ok(()); + } + ensure_open(state.terminal, item.stream_id)?; + if item.sequence != state.next_offset { + return Err(format!( + "input stream {} expected sequence {}, got {}", + item.stream_id, state.next_offset, item.sequence + )); + } + let next_offset = state + .next_offset + .checked_add(logical_item_count) + .ok_or_else(|| format!("input stream {} offset overflow", item.stream_id))?; + state.pending_acks.push_back(PendingAck { + sequence: item.sequence, + logical_item_count, + }); + state.next_offset = next_offset; + self.insert_input_streams(discovered); + Ok(()) + } + + fn validate_ack(&mut self, ack: &InputStreamAck) -> Result<(), String> { + let state = self + .inputs + .get_mut(&ack.stream_id) + .ok_or_else(|| format!("input stream {} is unknown", ack.stream_id))?; + let expected = state.pending_acks.front().ok_or_else(|| { + format!( + "input stream {} has no item awaiting acknowledgement", + ack.stream_id + ) + })?; + if ack.sequence != expected.sequence + || ack.logical_item_count != expected.logical_item_count + { + return Err(format!( + "input stream {} expected acknowledgement ({}, {}), got ({}, {})", + ack.stream_id, + expected.sequence, + expected.logical_item_count, + ack.sequence, + ack.logical_item_count + )); + } + state.pending_acks.pop_front(); + Ok(()) + } + + fn terminate_input(&mut self, stream_id: u64, offset: u64) -> Result<(), String> { + let state = self + .inputs + .get_mut(&stream_id) + .ok_or_else(|| format!("input stream {stream_id} is unknown"))?; + if let Some(discard_next_offset) = state.discard_next_offset { + if offset != discard_next_offset { + return Err(format!( + "cancelled input stream {stream_id} expected discarded terminal offset {discard_next_offset}, got {offset}" + )); + } + state.discard_next_offset = None; + return Ok(()); + } + ensure_open(state.terminal, stream_id)?; + if offset != state.next_offset { + return Err(format!( + "input stream {stream_id} expected terminal offset {}, got {offset}", + state.next_offset + )); + } + state.terminal = true; + Ok(()) + } + + fn cancel_input( + &mut self, + stream_id: u64, + offset: u64, + discard_in_flight: bool, + ) -> Result<(), String> { + let state = self + .inputs + .get_mut(&stream_id) + .ok_or_else(|| format!("input stream {stream_id} is unknown"))?; + ensure_open(state.terminal, stream_id)?; + let accepted_offset = state + .pending_acks + .front() + .map(|pending| pending.sequence) + .unwrap_or(state.next_offset); + if offset != accepted_offset { + return Err(format!( + "input stream {stream_id} expected cancellation offset {accepted_offset}, got {offset}" + )); + } + if discard_in_flight { + state.discard_next_offset = Some(state.next_offset); + } + state.next_offset = offset; + state.pending_acks.clear(); + state.terminal = true; + Ok(()) + } + + fn terminate_output(&mut self, stream_id: u64, offset: u64) -> Result<(), String> { + let state = self + .outputs + .get_mut(&stream_id) + .ok_or_else(|| format!("output stream {stream_id} is unknown"))?; + ensure_open(state.terminal, stream_id)?; + if offset != state.next_offset { + return Err(format!( + "output stream {stream_id} expected terminal offset {}, got {offset}", + state.next_offset + )); + } + state.terminal = true; + Ok(()) + } + + fn request_output_cancellation( + &mut self, + stream_id: u64, + offset: u64, + accept_terminal: bool, + ) -> Result<(), String> { + let state = self + .outputs + .get_mut(&stream_id) + .ok_or_else(|| format!("output stream {stream_id} is unknown"))?; + if !accept_terminal { + ensure_open(state.terminal, stream_id)?; + } + if state.cancellation_requested.is_some() { + return Err(format!( + "output stream {stream_id} already has a pending consumer cancellation" + )); + } + if offset > state.next_offset { + return Err(format!( + "output stream {stream_id} cannot cancel at future offset {offset}; latest observed offset is {}", + state.next_offset + )); + } + state.cancellation_requested = Some(offset); + Ok(()) + } + + fn confirm_output_cancellation(&mut self, stream_id: u64, offset: u64) -> Result<(), String> { + let requested = { + let state = self + .outputs + .get(&stream_id) + .ok_or_else(|| format!("output stream {stream_id} is unknown"))?; + ensure_open(state.terminal, stream_id)?; + state.cancellation_requested + }; + match requested { + Some(requested_offset) if offset != requested_offset => Err(format!( + "output stream {stream_id} expected cancellation confirmation at offset {requested_offset}, got {offset}" + )), + Some(_) => { + self.outputs + .get_mut(&stream_id) + .expect("output stream disappeared while confirming cancellation") + .terminal = true; + Ok(()) + } + None => self.terminate_output(stream_id, offset), + } + } + + fn finish(&mut self, finished: &InvocationSessionCompletion) -> Result<(), String> { + match finished.outcome.as_ref() { + Some(invocation_session_completion::Outcome::Success(_)) if !self.has_result => { + return Err( + "invocation completed successfully before publishing a result".to_string(), + ); + } + Some(invocation_session_completion::Outcome::Failure(failure)) => { + let kind = InvocationFailureKind::try_from(failure.kind) + .map_err(|_| format!("invalid invocation failure kind {}", failure.kind))?; + if kind == InvocationFailureKind::Unspecified { + return Err("invocation failure kind is unspecified".to_string()); + } + if failure.worker_error.is_some() && kind != InvocationFailureKind::Execution { + return Err( + "worker execution details require an execution failure kind".to_string() + ); + } + } + Some(invocation_session_completion::Outcome::Success(_)) => {} + None => return Err("invocation completion has no outcome".to_string()), + } + let unterminated_inputs = self + .inputs + .iter() + .filter_map(|(stream_id, state)| (!state.terminal).then_some(*stream_id)) + .collect::>(); + if !unterminated_inputs.is_empty() { + return Err(format!( + "invocation completed before input streams terminated: {unterminated_inputs:?}" + )); + } + let unacknowledged_inputs = self + .inputs + .iter() + .filter_map(|(stream_id, state)| (!state.pending_acks.is_empty()).then_some(*stream_id)) + .collect::>(); + if !unacknowledged_inputs.is_empty() { + return Err(format!( + "invocation completed with unacknowledged input streams: {unacknowledged_inputs:?}" + )); + } + let unterminated_outputs = self + .outputs + .iter() + .filter_map(|(stream_id, state)| (!state.terminal).then_some(*stream_id)) + .collect::>(); + if !unterminated_outputs.is_empty() { + return Err(format!( + "invocation completed before output streams terminated: {unterminated_outputs:?}" + )); + } + self.phase = SessionPhase::Complete; + Ok(()) + } + + fn validate_idempotency_key(&self, key: &Option) -> Result<(), String> { + let actual = required_idempotency_key(key)?; + let expected = self + .idempotency_key + .as_deref() + .ok_or_else(|| "invocation has no idempotency key".to_string())?; + if actual != expected { + return Err(format!( + "invocation idempotency key mismatch: expected {expected}, got {actual}" + )); + } + Ok(()) + } + + fn validate_identity( + &self, + agent_id: &Option, + idempotency_key: &Option, + ) -> Result<(), String> { + self.validate_idempotency_key(idempotency_key)?; + let actual = agent_id + .as_ref() + .ok_or_else(|| "invocation result has no agent identity".to_string())?; + let expected = self + .accepted_agent_id + .as_ref() + .ok_or_else(|| "invocation has no accepted agent identity".to_string())?; + if actual != expected { + return Err("invocation result agent identity differs from acceptance".to_string()); + } + Ok(()) + } + + fn ensure_new_input_streams(&self, stream_ids: &[u64]) -> Result<(), String> { + for stream_id in stream_ids { + if self.inputs.contains_key(stream_id) || self.outputs.contains_key(stream_id) { + return Err(format!("stream {stream_id} is already registered")); + } + } + Ok(()) + } + + fn insert_input_streams(&mut self, stream_ids: Vec) { + for stream_id in stream_ids { + self.inputs.insert(stream_id, InputState::default()); + } + } + + fn ensure_new_output_streams(&self, stream_ids: &[u64]) -> Result<(), String> { + for stream_id in stream_ids { + if self.inputs.contains_key(stream_id) || self.outputs.contains_key(stream_id) { + return Err(format!("stream {stream_id} is already registered")); + } + } + Ok(()) + } + + fn insert_output_streams(&mut self, stream_ids: Vec) { + for stream_id in stream_ids { + self.outputs.insert(stream_id, OutputState::default()); + } + } +} + +fn required_idempotency_key(key: &Option) -> Result<&str, String> { + let value = key + .as_ref() + .ok_or_else(|| "invocation has no idempotency key".to_string())? + .value + .as_str(); + if value.is_empty() { + Err("invocation idempotency key is empty".to_string()) + } else { + Ok(value) + } +} + +fn validate_cancel(cancel: &StreamCancel) -> Result<(), String> { + let role = StreamCancelRole::try_from(cancel.role) + .map_err(|_| format!("invalid stream cancellation role {}", cancel.role))?; + if role == StreamCancelRole::Unspecified { + return Err("stream cancellation role is unspecified".to_string()); + } + let reason = StreamCancelReason::try_from(cancel.reason) + .map_err(|_| format!("invalid stream cancellation reason {}", cancel.reason))?; + if reason == StreamCancelReason::Unspecified { + return Err("stream cancellation reason is unspecified".to_string()); + } + Ok(()) +} + +fn ensure_open(terminal: bool, stream_id: u64) -> Result<(), String> { + if terminal { + Err(format!( + "stream {stream_id} received an event after its terminal" + )) + } else { + Ok(()) + } +} + +fn stream_references(value: &SchemaValue) -> Result, String> { + fn visit( + value: &SchemaValue, + stream_ids: &mut Vec, + unique: &mut HashSet, + ) -> Result<(), String> { + match value + .value + .as_ref() + .ok_or_else(|| "schema value has no payload".to_string())? + { + schema_value::Value::RecordValue(record) => { + for field in &record.fields { + visit(field, stream_ids, unique)?; + } + } + schema_value::Value::VariantValue(variant) => { + if let Some(payload) = variant.payload.as_deref() { + visit(payload, stream_ids, unique)?; + } + } + schema_value::Value::TupleValue(tuple) => { + for element in &tuple.elements { + visit(element, stream_ids, unique)?; + } + } + schema_value::Value::ListValue(list) => { + for element in &list.elements { + visit(element, stream_ids, unique)?; + } + } + schema_value::Value::FixedListValue(list) => { + for element in &list.elements { + visit(element, stream_ids, unique)?; + } + } + schema_value::Value::MapValue(map) => { + for entry in &map.entries { + if let Some(key) = entry.key.as_ref() { + visit(key, stream_ids, unique)?; + } + if let Some(value) = entry.value.as_ref() { + visit(value, stream_ids, unique)?; + } + } + } + schema_value::Value::OptionValue(option) => { + if let Some(inner) = option.inner.as_deref() { + visit(inner, stream_ids, unique)?; + } + } + schema_value::Value::ResultValue(result) => match result.result.as_ref() { + Some(result_value::Result::Ok(value) | result_value::Result::Err(value)) => { + visit(value, stream_ids, unique)?; + } + Some(result_value::Result::OkUnit(_) | result_value::Result::ErrUnit(_)) => {} + None => return Err("result schema value has no payload".to_string()), + }, + schema_value::Value::UnionValue(union) => { + if let Some(body) = union.body.as_deref() { + visit(body, stream_ids, unique)?; + } + } + schema_value::Value::StreamReference(reference) => { + if !unique.insert(reference.stream_id) { + return Err(format!( + "stream {} is referenced more than once", + reference.stream_id + )); + } + stream_ids.push(reference.stream_id); + } + _ => {} + } + Ok(()) + } + + let mut stream_ids = Vec::new(); + visit(value, &mut stream_ids, &mut HashSet::new())?; + Ok(stream_ids) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::proto::golem::common::Empty; + use crate::proto::golem::schema::{RecordValue, SchemaValueStreamReference}; + use crate::proto::golem::worker::{ + InvocationFailure, InvocationRejected, InvocationStart, OutputStreamEnd, OutputStreamError, + OutputStreamItem, PublicInvocationStart, ResumeAttach, + }; + use prost::Message; + use test_r::test; + + const KEY: &str = "session-key"; + + fn key() -> Option { + Some(IdempotencyKey { + value: KEY.to_string(), + }) + } + + fn scalar(value: u32) -> SchemaValue { + SchemaValue { + value: Some(schema_value::Value::U8Value(value)), + } + } + + fn stream(stream_id: u64) -> SchemaValue { + SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id }, + )), + } + } + + fn record(fields: Vec) -> SchemaValue { + SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { fields })), + } + } + + fn public_request(request: public_invocation_request::Request) -> PublicInvocationRequest { + PublicInvocationRequest { + request: Some(request), + } + } + + fn trusted_request(request: invocation_request::Request) -> InvocationRequest { + InvocationRequest { + request: Some(request), + } + } + + fn response(response: invocation_response::Response) -> InvocationResponse { + InvocationResponse { + response: Some(response), + } + } + + fn public_start(input: SchemaValue) -> PublicInvocationRequest { + public_request(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: "app".to_string(), + environment_name: "env".to_string(), + agent_type_name: "agent-type".to_string(), + constructor_parameters: Some(record(Vec::new())), + method_name: "run".to_string(), + method_parameters: Some(input), + idempotency_key: key(), + ..Default::default() + }, + )) + } + + fn trusted_start(input: SchemaValue) -> InvocationRequest { + trusted_request(invocation_request::Request::Start(InvocationStart { + input: Some(input), + idempotency_key: key(), + ..Default::default() + })) + } + + fn agent_id() -> AgentId { + AgentId { + component_id: None, + name: "agent".to_string(), + } + } + + fn accepted() -> InvocationResponse { + response(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: Some(agent_id()), + idempotency_key: key(), + component_revision: Some(12), + }, + )) + } + + fn result(value: SchemaValue) -> InvocationResponse { + response(invocation_response::Response::Result( + InvocationSessionResult { + result: Some(invocation_session_result::Result::MethodResult(value)), + component_revision: Some(12), + agent_id: Some(agent_id()), + idempotency_key: key(), + ..Default::default() + }, + )) + } + + fn success() -> InvocationResponse { + response(invocation_response::Response::Finished( + InvocationSessionCompletion { + outcome: Some(invocation_session_completion::Outcome::Success(Empty {})), + }, + )) + } + + fn failure(kind: InvocationFailureKind) -> InvocationResponse { + response(invocation_response::Response::Finished( + InvocationSessionCompletion { + outcome: Some(invocation_session_completion::Outcome::Failure( + InvocationFailure { + kind: kind as i32, + code: "failed".to_string(), + message: "invocation failed".to_string(), + worker_error: (kind == InvocationFailureKind::Execution) + .then(Default::default), + }, + )), + }, + )) + } + + fn input_item(sequence: u64, payload: Payload) -> PublicInvocationRequest { + public_request(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id: 7, + sequence, + payload: Some(payload), + }, + )) + } + + fn cancel(stream_id: u64, role: StreamCancelRole, offset: u64) -> StreamCancel { + StreamCancel { + stream_id, + offset, + role: role as i32, + reason: StreamCancelReason::Cancelled as i32, + details: None, + } + } + + #[test] + fn legal_public_session_tracks_recursive_streams_acks_and_finish() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(record(vec![stream(7)]))) + .unwrap(); + assert!( + state + .validate_public_request(&input_item(0, Payload::PackedU8(vec![1, 2]))) + .is_err() + ); + state.validate_response(&accepted()).unwrap(); + state + .validate_public_request(&input_item(0, Payload::PackedU8(vec![1, 2]))) + .unwrap(); + state + .validate_response(&response(invocation_response::Response::InputAck( + InputStreamAck { + stream_id: 7, + sequence: 0, + logical_item_count: 2, + }, + ))) + .unwrap(); + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 7, + offset: 2, + }), + )) + .unwrap(); + state + .validate_response(&result(record(vec![stream(9), stream(10)]))) + .unwrap(); + state + .validate_response(&response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 9, + offset: 0, + value: Some(record(vec![stream(11)])), + }, + ))) + .unwrap(); + for (stream_id, offset) in [(9, 1), (10, 0), (11, 0)] { + state + .validate_response(&response(invocation_response::Response::OutputEnd( + OutputStreamEnd { stream_id, offset }, + ))) + .unwrap(); + } + state.validate_response(&success()).unwrap(); + assert!(state.is_complete()); + assert!(state.validate_response(&result(scalar(1))).is_err()); + } + + #[test] + fn legal_trusted_stream_free_session_is_accepted_result_finished() { + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&trusted_start(record(Vec::new()))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + state.validate_response(&result(scalar(1))).unwrap(); + state.validate_response(&success()).unwrap(); + } + + #[test] + fn resume_attach_requires_terminal_resume_unsupported_rejection() { + let resume = public_request(public_invocation_request::Request::ResumeAttach( + ResumeAttach { + idempotency_key: key(), + }, + )); + let mut state = InvocationSessionState::default(); + state.validate_public_request(&resume).unwrap(); + assert!(state.validate_response(&accepted()).is_err()); + assert!( + state + .validate_response(&response(invocation_response::Response::Rejected( + InvocationRejected { + reason: InvocationRejectionReason::Validation as i32, + idempotency_key: key(), + ..Default::default() + }, + ))) + .is_err() + ); + state + .validate_response(&response(invocation_response::Response::Rejected( + InvocationRejected { + reason: InvocationRejectionReason::ResumeUnsupported as i32, + idempotency_key: key(), + ..Default::default() + }, + ))) + .unwrap(); + assert!(state.is_complete()); + assert!(state.validate_public_request(&resume).is_err()); + } + + #[test] + fn rejection_is_terminal_for_both_directions() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(record(Vec::new()))) + .unwrap(); + state + .validate_response(&response(invocation_response::Response::Rejected( + InvocationRejected { + reason: InvocationRejectionReason::Validation as i32, + idempotency_key: key(), + ..Default::default() + }, + ))) + .unwrap(); + assert!(state.validate_response(&accepted()).is_err()); + assert!( + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd::default()), + )) + .is_err() + ); + } + + #[test] + fn unknown_streams_ack_mismatch_and_post_terminal_events_are_rejected() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(stream(7))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + assert!( + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 8, + offset: 0, + }), + )) + .is_err() + ); + state + .validate_public_request(&input_item(0, Payload::PackedU8(vec![1, 2, 3]))) + .unwrap(); + assert!( + state + .validate_response(&response(invocation_response::Response::InputAck( + InputStreamAck { + stream_id: 7, + sequence: 0, + logical_item_count: 2, + }, + ))) + .is_err() + ); + state + .validate_response(&response(invocation_response::Response::InputAck( + InputStreamAck { + stream_id: 7, + sequence: 0, + logical_item_count: 3, + }, + ))) + .unwrap(); + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 7, + offset: 3, + }), + )) + .unwrap(); + assert!( + state + .validate_public_request(&input_item(3, Payload::Value(scalar(1)))) + .is_err() + ); + } + + #[test] + fn received_output_cancellation_may_race_with_stream_terminal() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(record(Vec::new()))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + state.validate_response(&result(stream(9))).unwrap(); + state + .validate_response(&response(invocation_response::Response::OutputEnd( + OutputStreamEnd { + stream_id: 9, + offset: 0, + }, + ))) + .unwrap(); + + assert!( + state + .validate_public_request(&public_request( + public_invocation_request::Request::StreamCancel(cancel( + 9, + StreamCancelRole::OutputConsumer, + 0, + )), + )) + .is_err(), + "an output-consumer cancellation must not be accepted after the stream terminal" + ); + + let cancellation = public_request(public_invocation_request::Request::StreamCancel( + cancel(9, StreamCancelRole::OutputConsumer, 0), + )); + state + .validate_received_public_request(&cancellation) + .unwrap(); + assert!( + state + .validate_received_public_request(&cancellation) + .is_err() + ); + } + + #[test] + fn input_items_register_recursively_nested_streams() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(stream(7))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + state + .validate_public_request(&input_item(0, Payload::Value(record(vec![stream(8)])))) + .unwrap(); + state + .validate_response(&response(invocation_response::Response::InputAck( + InputStreamAck { + stream_id: 7, + sequence: 0, + logical_item_count: 1, + }, + ))) + .unwrap(); + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 7, + offset: 1, + }), + )) + .unwrap(); + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 8, + offset: 0, + }), + )) + .unwrap(); + state.validate_response(&result(scalar(1))).unwrap(); + state.validate_response(&success()).unwrap(); + } + + #[test] + fn all_role_appropriate_cancellations_are_unique_terminals() { + let mut input = InvocationSessionState::default(); + input + .validate_public_request(&public_start(stream(7))) + .unwrap(); + input.validate_response(&accepted()).unwrap(); + let input_cancel = public_request(public_invocation_request::Request::StreamCancel( + cancel(7, StreamCancelRole::InputProducer, 0), + )); + input.validate_public_request(&input_cancel).unwrap(); + assert!(input.validate_public_request(&input_cancel).is_err()); + + let mut input_consumer = InvocationSessionState::default(); + input_consumer + .validate_public_request(&public_start(stream(7))) + .unwrap(); + input_consumer.validate_response(&accepted()).unwrap(); + input_consumer + .validate_response(&response(invocation_response::Response::StreamCancel( + cancel(7, StreamCancelRole::InputConsumer, 0), + ))) + .unwrap(); + + let mut output = InvocationSessionState::default(); + output + .validate_public_request(&public_start(record(Vec::new()))) + .unwrap(); + output.validate_response(&accepted()).unwrap(); + output.validate_response(&result(stream(9))).unwrap(); + output + .validate_public_request(&public_request( + public_invocation_request::Request::StreamCancel(cancel( + 9, + StreamCancelRole::OutputConsumer, + 0, + )), + )) + .unwrap(); + assert!( + output + .validate_public_request(&public_request( + public_invocation_request::Request::StreamCancel(cancel( + 9, + StreamCancelRole::OutputConsumer, + 0, + )), + )) + .is_err() + ); + output + .validate_response(&response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 9, + offset: 0, + value: Some(scalar(1)), + }, + ))) + .unwrap(); + output + .validate_response(&response(invocation_response::Response::StreamCancel( + cancel(9, StreamCancelRole::OutputProducer, 0), + ))) + .unwrap(); + output.validate_response(&success()).unwrap(); + + let mut output_producer = InvocationSessionState::default(); + output_producer + .validate_public_request(&public_start(record(Vec::new()))) + .unwrap(); + output_producer.validate_response(&accepted()).unwrap(); + output_producer + .validate_response(&result(stream(9))) + .unwrap(); + output_producer + .validate_response(&response(invocation_response::Response::StreamCancel( + cancel(9, StreamCancelRole::OutputProducer, 0), + ))) + .unwrap(); + } + + #[test] + fn input_consumer_cancellation_abandons_an_unacknowledged_item() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(stream(7))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + state + .validate_public_request(&input_item(0, Payload::Value(scalar(1)))) + .unwrap(); + state + .validate_response(&response(invocation_response::Response::StreamCancel( + cancel(7, StreamCancelRole::InputConsumer, 0), + ))) + .unwrap(); + state + .validate_public_request(&input_item(1, Payload::Value(scalar(2)))) + .unwrap(); + state + .validate_public_request(&public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 7, + offset: 2, + }), + )) + .unwrap(); + state.validate_response(&result(scalar(2))).unwrap(); + state.validate_response(&success()).unwrap(); + } + + #[test] + fn duplicate_input_end_after_consumer_cancellation_is_rejected() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(stream(7))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + state + .validate_response(&response(invocation_response::Response::StreamCancel( + cancel(7, StreamCancelRole::InputConsumer, 0), + ))) + .unwrap(); + let end = public_request(public_invocation_request::Request::InputEnd( + InputStreamEnd { + stream_id: 7, + offset: 0, + }, + )); + state.validate_public_request(&end).unwrap(); + + assert!( + state.validate_public_request(&end).is_err(), + "an input stream must not accept the producer terminal more than once" + ); + } + + #[test] + fn recursive_registration_is_transactional_and_attachment_revocation_is_illegal() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(record(Vec::new()))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + assert!( + state + .validate_response(&result(record(vec![stream(9), stream(9)]))) + .is_err() + ); + state.validate_response(&result(stream(9))).unwrap(); + assert!( + state + .validate_response(&response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 9, + offset: 0, + value: Some(stream(9)), + }, + ))) + .is_err() + ); + state + .validate_response(&response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 9, + offset: 0, + value: Some(scalar(1)), + }, + ))) + .unwrap(); + assert!( + state + .validate_response(&response(invocation_response::Response::AttachmentRevoked( + Default::default() + ),)) + .is_err() + ); + } + + #[test] + fn stream_ids_are_unique_across_input_and_output_directions() { + let mut input_first = InvocationSessionState::default(); + input_first + .validate_public_request(&public_start(stream(7))) + .unwrap(); + input_first.validate_response(&accepted()).unwrap(); + assert!(input_first.validate_response(&result(stream(7))).is_err()); + + let mut output_first = InvocationSessionState::default(); + output_first + .validate_public_request(&public_start(stream(7))) + .unwrap(); + output_first.validate_response(&accepted()).unwrap(); + output_first.validate_response(&result(stream(9))).unwrap(); + assert!( + output_first + .validate_public_request(&input_item(0, Payload::Value(stream(9)))) + .is_err() + ); + } + + #[test] + fn finish_requires_result_terminals_and_safe_failure_kind() { + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&public_start(record(Vec::new()))) + .unwrap(); + state.validate_response(&accepted()).unwrap(); + assert!(state.validate_response(&success()).is_err()); + assert!( + state + .validate_response(&failure(InvocationFailureKind::Unspecified)) + .is_err() + ); + let mut protocol_with_worker_error = failure(InvocationFailureKind::Protocol); + let Some(invocation_response::Response::Finished(InvocationSessionCompletion { + outcome: Some(invocation_session_completion::Outcome::Failure(protocol_failure)), + })) = protocol_with_worker_error.response.as_mut() + else { + unreachable!("failure helper returned the wrong response") + }; + protocol_failure.worker_error = Some(Default::default()); + assert!( + state + .validate_response(&protocol_with_worker_error) + .is_err() + ); + state + .validate_response(&failure(InvocationFailureKind::Execution)) + .unwrap(); + } + + fn round_trip(message: M) + where + M: Message + Default + PartialEq + std::fmt::Debug, + { + let encoded = message.encode_to_vec(); + assert_eq!(M::decode(encoded.as_slice()).unwrap(), message); + } + + #[test] + fn public_and_internal_envelopes_round_trip_all_gol_91_variants() { + round_trip(public_start(record(Vec::new()))); + round_trip(trusted_start(record(Vec::new()))); + round_trip(public_request( + public_invocation_request::Request::ResumeAttach(ResumeAttach { + idempotency_key: key(), + }), + )); + round_trip(input_item(4, Payload::PackedU8(vec![1, 2, 3]))); + round_trip(public_request( + public_invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 7, + offset: 4, + }), + )); + for role in [ + StreamCancelRole::InputProducer, + StreamCancelRole::InputConsumer, + StreamCancelRole::OutputProducer, + StreamCancelRole::OutputConsumer, + ] { + round_trip(StreamCancel { + stream_id: 7, + offset: 8, + role: role as i32, + reason: StreamCancelReason::Protocol as i32, + details: Some("cancelled".to_string()), + }); + } + round_trip(accepted()); + round_trip(response(invocation_response::Response::Rejected( + InvocationRejected { + reason: InvocationRejectionReason::ResumeUnsupported as i32, + idempotency_key: key(), + ..Default::default() + }, + ))); + round_trip(result(stream(9))); + round_trip(response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 9, + offset: 0, + value: Some(scalar(1)), + }, + ))); + round_trip(response(invocation_response::Response::OutputEnd( + OutputStreamEnd { + stream_id: 9, + offset: 1, + }, + ))); + round_trip(response(invocation_response::Response::OutputError( + OutputStreamError { + stream_id: 9, + offset: 1, + details: "failed".to_string(), + }, + ))); + round_trip(response(invocation_response::Response::InputAck( + InputStreamAck { + stream_id: 7, + sequence: 0, + logical_item_count: 3, + }, + ))); + round_trip(response(invocation_response::Response::StreamCancel( + cancel(9, StreamCancelRole::OutputProducer, 1), + ))); + round_trip(response(invocation_response::Response::AttachmentRevoked( + Default::default(), + ))); + round_trip(success()); + round_trip(failure(InvocationFailureKind::Protocol)); + round_trip(failure(InvocationFailureKind::Execution)); + } +} diff --git a/golem-api-grpc/src/lib.rs b/golem-api-grpc/src/lib.rs index d59f80deed..87f36c5607 100644 --- a/golem-api-grpc/src/lib.rs +++ b/golem-api-grpc/src/lib.rs @@ -15,6 +15,8 @@ #[cfg(test)] test_r::enable!(); +pub mod invocation_session_protocol; + #[allow(clippy::large_enum_variant)] pub mod proto { use crate::proto::golem::worker::UpdateMode; diff --git a/golem-common/src/model/agent/extraction.rs b/golem-common/src/model/agent/extraction.rs index a33130a1bd..9b0df03378 100644 --- a/golem-common/src/model/agent/extraction.rs +++ b/golem-common/src/model/agent/extraction.rs @@ -21,6 +21,7 @@ use crate::schema::tool::validation::validate_tool; use crate::schema::tool::wit::wire as tool_wire; use crate::wasmtime_config::create_wasmtime_config; use anyhow::anyhow; +use golem_schema::schema::SchemaValueStreamHandleRep; use golem_schema::schema::wit::{ QuotaTokenHandleDropper, QuotaTokenHandleRep, SecretHandleDropper, SecretHandleRep, }; @@ -530,6 +531,10 @@ fn is_secret_resource(interface_name: &str, resource_name: &str) -> bool { ) } +fn is_schema_value_stream_resource(interface_name: &str, resource_name: &str) -> bool { + interface_name == "golem:core/types@2.0.0" && resource_name == "schema-value-stream" +} + fn dynamic_import( name: &str, engine: &Engine, @@ -596,6 +601,12 @@ fn dynamic_import( ResourceType::host::(), |_store, _rep| Ok(()), )?; + } else if is_schema_value_stream_resource(&name, &inner_name) { + instance.resource( + &inner_name, + ResourceType::host::(), + |_store, _rep| Ok(()), + )?; } else if &inner_name != "pollable" && inner_name != "wasi-io-pollable" && &inner_name != "input-stream" diff --git a/golem-common/src/model/component_metadata.rs b/golem-common/src/model/component_metadata.rs index bd77b0528d..b8dcc01b59 100644 --- a/golem-common/src/model/component_metadata.rs +++ b/golem-common/src/model/component_metadata.rs @@ -795,17 +795,21 @@ mod protobuf { } } - impl From for golem_api_grpc::proto::golem::component::ComponentMetadata { - fn from(value: ComponentMetadata) -> Self { - value.data.as_ref().clone().into() + impl TryFrom for golem_api_grpc::proto::golem::component::ComponentMetadata { + type Error = String; + + fn try_from(value: ComponentMetadata) -> Result { + value.data.as_ref().clone().try_into() } } - impl From + impl TryFrom for golem_api_grpc::proto::golem::component::ComponentMetadata { - fn from(value: ComponentMetadataInnerData) -> Self { - Self { + type Error = String; + + fn try_from(value: ComponentMetadataInnerData) -> Result { + Ok(Self { known_exports: Some(value.known_exports.into()), producers: value .producers @@ -823,21 +827,14 @@ mod protobuf { agent_type_provision_configs: value .agent_type_provision_configs .into_iter() - .map(|(k, v)| { - ( - k.0, - golem_api_grpc::proto::golem::component::AgentTypeProvisionConfig::from( - v, - ), - ) - }) - .collect(), + .map(|(k, v)| v.try_into().map(|config| (k.0, config))) + .collect::>()?, tools: value .tools .into_iter() .map(|(name, metadata)| (name.into_inner(), metadata.into())) .collect(), - } + }) } } @@ -1049,21 +1046,23 @@ mod protobuf { } } - impl From + impl TryFrom for golem_api_grpc::proto::golem::component::AgentTypeProvisionConfig { - fn from(config: AgentTypeProvisionConfig) -> Self { + type Error = String; + + fn try_from(config: AgentTypeProvisionConfig) -> Result { use crate::base_model::component::{InitialAgentFile, InstalledPlugin}; - Self { + Ok(Self { initial_permissions: crate::serialization::serialize(&config.initial_permissions) .expect("failed to serialize agent initial permission card"), env: config.env.into_iter().collect(), config: config .config .into_iter() - .map(golem_api_grpc::proto::golem::worker::TypedAgentConfigEntry::from) - .collect(), + .map(TryInto::try_into) + .collect::>()?, plugins: config .plugins .into_iter() @@ -1078,7 +1077,7 @@ mod protobuf { golem_api_grpc::proto::golem::component::InitialAgentFile::from(f) }) .collect(), - } + }) } } @@ -1455,7 +1454,8 @@ mod tests { )]), ); - let proto: golem_api_grpc::proto::golem::component::ComponentMetadata = metadata.into(); + let proto: golem_api_grpc::proto::golem::component::ComponentMetadata = + metadata.try_into().unwrap(); let decoded = ComponentMetadata::try_from(proto).unwrap(); assert_eq!( @@ -1479,7 +1479,8 @@ mod tests { BTreeMap::new(), ); - let proto: golem_api_grpc::proto::golem::component::ComponentMetadata = metadata.into(); + let proto: golem_api_grpc::proto::golem::component::ComponentMetadata = + metadata.try_into().unwrap(); let decoded = ComponentMetadata::try_from(proto).unwrap(); assert!(decoded.memories()[0].shared); @@ -1556,7 +1557,7 @@ mod tests { fn component_metadata_grpc_roundtrip_preserves_tool_envelope() { let metadata = metadata_with_tool(); let proto: golem_api_grpc::proto::golem::component::ComponentMetadata = - metadata.clone().into(); + metadata.clone().try_into().unwrap(); let decoded = ComponentMetadata::try_from(proto).unwrap(); assert_eq!(decoded, metadata); diff --git a/golem-common/src/model/oplog/matcher.rs b/golem-common/src/model/oplog/matcher.rs index 1d9e4b9f18..2e6d6d5f04 100644 --- a/golem-common/src/model/oplog/matcher.rs +++ b/golem-common/src/model/oplog/matcher.rs @@ -783,6 +783,7 @@ impl PublicOplogEntry { SchemaValue::QuotaToken(payload) => { Self::string_match(&payload.resource_name, path_stack, query_path, query) } + SchemaValue::Stream(_) => false, } } } diff --git a/golem-common/src/model/oplog/protobuf.rs b/golem-common/src/model/oplog/protobuf.rs index 575eaccbc8..e00cca8663 100644 --- a/golem-common/src/model/oplog/protobuf.rs +++ b/golem-common/src/model/oplog/protobuf.rs @@ -70,14 +70,16 @@ use golem_api_grpc::proto::golem::worker::{ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::num::NonZeroU64; -impl From +impl TryFrom for golem_api_grpc::proto::golem::worker::PublicTypedAgentConfigEntry { - fn from(value: PublicTypedAgentConfigEntry) -> Self { - Self { + type Error = String; + + fn try_from(value: PublicTypedAgentConfigEntry) -> Result { + Ok(Self { path: value.path, - value: Some(value.value.into()), - } + value: Some(value.value.try_into()?), + }) } } @@ -1028,8 +1030,8 @@ impl TryFrom for golem_api_grpc::proto::golem::worker::OplogEn config: create .local_agent_config .into_iter() - .map(Into::into) - .collect(), + .map(TryInto::try_into) + .collect::>()?, created_by: Some(create.created_by.into()), environment_id: Some(create.environment_id.into()), parent: create.parent.map(Into::into), @@ -1056,7 +1058,7 @@ impl TryFrom for golem_api_grpc::proto::golem::worker::OplogEn observational_owner: start .observational_owner .map(|id| id.as_u64()), - request: start.request.map(Into::into), + request: start.request.map(TryInto::try_into).transpose()?, durable_function_type: Some(start.durable_function_type.into()), }, )), @@ -1067,7 +1069,7 @@ impl TryFrom for golem_api_grpc::proto::golem::worker::OplogEn golem_api_grpc::proto::golem::worker::EndParameters { timestamp: Some(end.timestamp.into()), start_index: end.start_index.as_u64(), - response: end.response.map(Into::into), + response: end.response.map(TryInto::try_into).transpose()?, forced_commit: end.forced_commit, }, )), @@ -1078,7 +1080,7 @@ impl TryFrom for golem_api_grpc::proto::golem::worker::OplogEn golem_api_grpc::proto::golem::worker::CancelledParameters { timestamp: Some(cancelled.timestamp.into()), start_index: cancelled.start_index.as_u64(), - partial: cancelled.partial.map(Into::into), + partial: cancelled.partial.map(TryInto::try_into).transpose()?, }, )), } @@ -1587,7 +1589,7 @@ impl TryFrom for golem_api_grpc::proto::golem::worker::OplogEn timestamp: Some(params.timestamp.into()), parent_start_index: params.parent_start_index.as_u64(), kind: host_stream_kind_to_proto(params.kind) as i32, - payload: Some(params.payload.into()), + payload: Some(params.payload.try_into()?), }, )), } @@ -1838,7 +1840,7 @@ impl TryFrom Invocation::AgentInitialization( golem_api_grpc::proto::golem::worker::PublicAgentInitializationInvocation { idempotency_key: Some(init.idempotency_key.into()), - constructor_parameters: Some(init.constructor_parameters.into()), + constructor_parameters: Some(init.constructor_parameters.try_into()?), trace_id: init.trace_id.to_string(), trace_states: init.trace_states, invocation_context, @@ -1851,7 +1853,7 @@ impl TryFrom golem_api_grpc::proto::golem::worker::PublicAgentMethodInvocation { idempotency_key: Some(method.idempotency_key.into()), method_name: method.method_name, - function_input: Some(method.function_input.into()), + function_input: Some(method.function_input.try_into()?), trace_id: method.trace_id.to_string(), trace_states: method.trace_states, invocation_context, @@ -2046,10 +2048,10 @@ impl TryFrom use golem_api_grpc::proto::golem::worker::public_agent_invocation_result::Result as ProtoResult; let result = match value { PublicAgentInvocationResult::AgentInitialization(output) => { - ProtoResult::AgentInitializationOutput(output.output.into()) + ProtoResult::AgentInitializationOutput(output.output.try_into()?) } PublicAgentInvocationResult::AgentMethod(output) => { - ProtoResult::AgentMethodOutput(output.output.into()) + ProtoResult::AgentMethodOutput(output.output.try_into()?) } PublicAgentInvocationResult::ManualUpdate(_) => { ProtoResult::ManualUpdate(golem_api_grpc::proto::golem::common::Empty {}) diff --git a/golem-common/src/model/worker.rs b/golem-common/src/model/worker.rs index 0589e1f90f..bb46fcb911 100644 --- a/golem-common/src/model/worker.rs +++ b/golem-common/src/model/worker.rs @@ -177,8 +177,10 @@ mod protobuf { } } - impl From for golem_api_grpc::proto::golem::worker::AgentMetadata { - fn from(value: AgentMetadataDto) -> Self { + impl TryFrom for golem_api_grpc::proto::golem::worker::AgentMetadata { + type Error = String; + + fn try_from(value: AgentMetadataDto) -> Result { let mut owned_resources = Vec::new(); for instance in value.exported_resource_instances { owned_resources.push(golem_api_grpc::proto::golem::worker::ResourceDescription { @@ -189,12 +191,16 @@ mod protobuf { }); } - Self { + Ok(Self { agent_id: Some(value.agent_id.into()), environment_id: Some(value.environment_id.into()), created_by: Some(value.created_by.into()), env: value.env, - config: value.config.into_iter().map(Into::into).collect(), + config: value + .config + .into_iter() + .map(TryInto::try_into) + .collect::>()?, status: value.status.into(), component_revision: value.component_revision.into(), retry_count: value.retry_count, @@ -222,7 +228,7 @@ mod protobuf { .collect(), oplog_idx: u64::from(value.last_oplog_index), fingerprint: Some(value.fingerprint.0.into()), - } + }) } } @@ -425,12 +431,16 @@ mod protobuf { } } - impl From for golem_api_grpc::proto::golem::worker::TypedAgentConfigEntry { - fn from(value: TypedAgentConfigEntry) -> Self { - Self { + impl TryFrom + for golem_api_grpc::proto::golem::worker::TypedAgentConfigEntry + { + type Error = String; + + fn try_from(value: TypedAgentConfigEntry) -> Result { + Ok(Self { path: value.path, - value: Some(value.value.into()), - } + value: Some(value.value.try_into()?), + }) } } } diff --git a/golem-common/src/schema/agent/mod.rs b/golem-common/src/schema/agent/mod.rs index 8bdea52faa..ed5eb45653 100644 --- a/golem-common/src/schema/agent/mod.rs +++ b/golem-common/src/schema/agent/mod.rs @@ -39,6 +39,7 @@ use crate::schema::graph::{SchemaGraph, TypedSchemaValue}; 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 golem_schema_derive::{FromSchema, IntoSchema}; use serde::{Deserialize, Serialize}; @@ -374,6 +375,29 @@ pub struct AgentMethodSchema { pub read_only: Option, } +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(|_| ()) + } + + /// Returns whether a caller-supplied input or the output of this method can + /// contain a stream, following references through the owning agent's graph. + pub fn uses_streams(&self, graph: &SchemaGraph) -> bool { + self.input_schema + .fields() + .iter() + .filter(|field| matches!(field.source, FieldSource::UserSupplied)) + .any(|field| contains_stream_in_graph(graph, &field.schema)) + || self + .output_schema + .schema() + .is_some_and(|output| contains_stream_in_graph(graph, output)) + } +} + /// Dependent agent type, schema-layer form. /// /// Owns its own [`SchemaGraph`] — a dependent agent is independently @@ -496,11 +520,8 @@ impl AgentTypeSchema { agent_types.into_iter().map(Self::normalized).collect() } - /// Validates the semantic constraints of the agent type. Mirrors the legacy - /// `AgentType::validate`: ephemeral agents must not declare read-only - /// methods (there is no shared state to read from). Additionally rejects - /// the WASI P3 stub types (`future`/`stream`) anywhere in the agent's - /// schemas — see [`reject_p3_stub_types`]. + /// Validates semantic constraints of the agent type, including stream + /// placement and definitions that are not reachable from an allowed use. pub fn validate(&self) -> Result<(), String> { if self.mode == AgentMode::Ephemeral { for method in &self.methods { @@ -514,142 +535,156 @@ impl AgentTypeSchema { } } } - reject_p3_stub_types(self) - } -} - -/// Rejects the WASI P3 stub types (`future`/`stream`) anywhere in an agent -/// type's schemas: shared type definitions, constructor and method inputs, -/// method outputs, config value types, and the same positions of every -/// dependency. -/// -/// These types parse ([`SchemaType::Future`] / [`SchemaType::Stream`]) but -/// have no [`SchemaValue`] representation and cannot be marshalled across the -/// invocation boundary, so accepting them at upload time would only defer the -/// failure to invocation time as a confusing shape mismatch. `error-context` -/// has no [`SchemaType`] representation at all, so it cannot occur here. -fn reject_p3_stub_types(agent_type: &AgentTypeSchema) -> Result<(), String> { - check_signatures_for_p3_stubs( - &agent_type.type_name, - None, - &agent_type.schema, - &agent_type.constructor, - &agent_type.methods, - )?; - for config in &agent_type.config { - if let Some(kind) = find_p3_stub(&config.value_type) { - return Err(p3_stub_error( - &agent_type.type_name, - &format!("config value at path '{}'", config.path.join(".")), - kind, - )); + validate_agent_type_placement(self).map_err(|errors| { + errors + .into_iter() + .map(|error| error.to_string()) + .collect::>() + .join("; ") + })?; + validate_schema_definitions(&self.type_name, None, &self.schema, &self.methods)?; + for dependency in &self.dependencies { + validate_schema_definitions( + &self.type_name, + Some(&dependency.type_name), + &dependency.schema, + &dependency.methods, + )?; } + Ok(()) } - for dependency in &agent_type.dependencies { - check_signatures_for_p3_stubs( - &agent_type.type_name, - Some(&dependency.type_name), - &dependency.schema, - &dependency.constructor, - &dependency.methods, - )?; - } - Ok(()) } -/// Checks one (graph, constructor, methods) signature set for P3 stub types. -/// `dependency` prefixes the reported location when the set belongs to an -/// [`AgentDependencySchema`] rather than the agent type itself. -fn check_signatures_for_p3_stubs( +fn validate_schema_definitions( agent: &AgentTypeName, dependency: Option<&str>, graph: &SchemaGraph, - constructor: &AgentConstructorSchema, methods: &[AgentMethodSchema], ) -> Result<(), String> { - let loc = |location: String| match dependency { - Some(dep) => format!("dependency '{dep}' {location}"), - None => location, - }; - for def in &graph.defs { - if let Some(kind) = find_p3_stub(&def.body) { - return Err(p3_stub_error( - agent, - &loc(format!("shared type definition '{}'", def.id)), - kind, - )); - } - } - for field in constructor.input_schema.fields() { - if let Some(kind) = find_p3_stub(&field.schema) { - return Err(p3_stub_error( - agent, - &loc(format!("constructor parameter '{}'", field.name)), - kind, + let owner = dependency + .map(|name| format!("dependency '{name}' ")) + .unwrap_or_default(); + let method_definitions = methods + .iter() + .flat_map(|method| { + let inputs = method + .input_schema + .fields() + .iter() + .filter(|field| matches!(field.source, FieldSource::UserSupplied)) + .map(|field| &field.schema); + let output = method.output_schema.schema().into_iter(); + inputs.chain(output) + }) + .flat_map(|root| reachable_defs(graph, root)) + .map(|definition| definition.id) + .collect::>(); + + for definition in &graph.defs { + if contains_future(&definition.body) { + return Err(format!( + "Agent type '{agent}' {owner}shared type definition '{}' contains the unsupported type 'future'", + definition.id )); } - } - for method in methods { - for field in method.input_schema.fields() { - if let Some(kind) = find_p3_stub(&field.schema) { - return Err(p3_stub_error( - agent, - &loc(format!( - "method '{}' parameter '{}'", - method.name, field.name - )), - kind, - )); - } - } - if let Some(output) = method.output_schema.schema() - && let Some(kind) = find_p3_stub(output) + if contains_stream_in_graph(graph, &definition.body) + && !method_definitions.contains(&definition.id) { - return Err(p3_stub_error( - agent, - &loc(format!("method '{}' output", method.name)), - kind, + return Err(format!( + "Agent type '{agent}' {owner}shared type definition '{}' contains a stream but is not reachable from a caller-supplied method input or method output", + definition.id )); } } Ok(()) } -fn p3_stub_error(agent: &AgentTypeName, location: &str, kind: &str) -> String { - format!( - "Agent type '{agent}' uses the unsupported type '{kind}' in {location}. \ - Stream, future, and error-context types have no value representation and \ - cannot be used in agent constructor, method, or config schemas." - ) +pub fn contains_stream_in_graph(graph: &SchemaGraph, ty: &SchemaType) -> bool { + contains_stream(ty) + || reachable_defs(graph, ty) + .iter() + .any(|definition| contains_stream(&definition.body)) +} + +fn contains_future(ty: &SchemaType) -> bool { + match ty { + SchemaType::Future { .. } => true, + SchemaType::Record { fields, .. } => { + fields.iter().any(|field| contains_future(&field.body)) + } + SchemaType::Variant { cases, .. } => cases + .iter() + .filter_map(|case| case.payload.as_ref()) + .any(contains_future), + SchemaType::Tuple { elements, .. } => elements.iter().any(contains_future), + SchemaType::List { element, .. } | SchemaType::FixedList { element, .. } => { + contains_future(element) + } + SchemaType::Map { key, value, .. } => contains_future(key) || contains_future(value), + SchemaType::Option { inner, .. } => contains_future(inner), + SchemaType::Result { spec, .. } => { + spec.ok.as_deref().is_some_and(contains_future) + || spec.err.as_deref().is_some_and(contains_future) + } + SchemaType::Union { spec, .. } => spec + .branches + .iter() + .any(|branch| contains_future(&branch.body)), + SchemaType::Secret { spec, .. } => contains_future(&spec.inner), + SchemaType::Stream { inner, .. } => inner.as_deref().is_some_and(contains_future), + SchemaType::Ref { .. } + | SchemaType::Bool { .. } + | SchemaType::S8 { .. } + | SchemaType::S16 { .. } + | SchemaType::S32 { .. } + | SchemaType::S64 { .. } + | SchemaType::U8 { .. } + | SchemaType::U16 { .. } + | SchemaType::U32 { .. } + | SchemaType::U64 { .. } + | SchemaType::F32 { .. } + | SchemaType::F64 { .. } + | SchemaType::Char { .. } + | SchemaType::String { .. } + | SchemaType::Enum { .. } + | SchemaType::Flags { .. } + | SchemaType::Text { .. } + | SchemaType::Binary { .. } + | SchemaType::Path { .. } + | SchemaType::Url { .. } + | SchemaType::Datetime { .. } + | SchemaType::Duration { .. } + | SchemaType::Quantity { .. } + | SchemaType::QuotaToken { .. } => false, + } } -/// Finds the first WASI P3 stub node (`future`/`stream`) in `ty`, descending -/// through all structural children but not following refs — named definitions -/// are scanned directly by [`check_signatures_for_p3_stubs`]. Returns the -/// offending node's kind name. -fn find_p3_stub(ty: &SchemaType) -> Option<&'static str> { +fn contains_stream(ty: &SchemaType) -> bool { match ty { - SchemaType::Future { .. } => Some("future"), - SchemaType::Stream { .. } => Some("stream"), - SchemaType::Record { fields, .. } => fields.iter().find_map(|f| find_p3_stub(&f.body)), + SchemaType::Stream { .. } => true, + SchemaType::Record { fields, .. } => { + fields.iter().any(|field| contains_stream(&field.body)) + } SchemaType::Variant { cases, .. } => cases .iter() - .find_map(|c| c.payload.as_ref().and_then(find_p3_stub)), - SchemaType::Tuple { elements, .. } => elements.iter().find_map(find_p3_stub), + .filter_map(|case| case.payload.as_ref()) + .any(contains_stream), + SchemaType::Tuple { elements, .. } => elements.iter().any(contains_stream), SchemaType::List { element, .. } | SchemaType::FixedList { element, .. } => { - find_p3_stub(element) + contains_stream(element) + } + SchemaType::Map { key, value, .. } => contains_stream(key) || contains_stream(value), + SchemaType::Option { inner, .. } => contains_stream(inner), + SchemaType::Result { spec, .. } => { + spec.ok.as_deref().is_some_and(contains_stream) + || spec.err.as_deref().is_some_and(contains_stream) } - SchemaType::Map { key, value, .. } => find_p3_stub(key).or_else(|| find_p3_stub(value)), - SchemaType::Option { inner, .. } => find_p3_stub(inner), - SchemaType::Result { spec, .. } => spec - .ok - .as_deref() - .and_then(find_p3_stub) - .or_else(|| spec.err.as_deref().and_then(find_p3_stub)), - SchemaType::Union { spec, .. } => spec.branches.iter().find_map(|b| find_p3_stub(&b.body)), - // Leaf nodes carrying no child `SchemaType`. Listed explicitly (no - // wildcard) so a future child-bearing variant forces this match to be - // updated. + SchemaType::Union { spec, .. } => spec + .branches + .iter() + .any(|branch| contains_stream(&branch.body)), + SchemaType::Secret { spec, .. } => contains_stream(&spec.inner), + SchemaType::Future { inner, .. } => inner.as_deref().is_some_and(contains_stream), SchemaType::Ref { .. } | SchemaType::Bool { .. } | SchemaType::S8 { .. } @@ -673,8 +708,7 @@ fn find_p3_stub(ty: &SchemaType) -> Option<&'static str> { | SchemaType::Datetime { .. } | SchemaType::Duration { .. } | SchemaType::Quantity { .. } - | SchemaType::Secret { .. } - | SchemaType::QuotaToken { .. } => None, + | SchemaType::QuotaToken { .. } => false, } } diff --git a/golem-common/src/schema/agent/tests.rs b/golem-common/src/schema/agent/tests.rs index d4f72f488d..1111a5e96c 100644 --- a/golem-common/src/schema/agent/tests.rs +++ b/golem-common/src/schema/agent/tests.rs @@ -17,13 +17,14 @@ use crate::base_model::agent::{AgentConfigSource, AgentMode, AgentTypeName, Snap use crate::schema::agent::{ AgentConfigDeclarationSchema, AgentConstructorSchema, AgentDependencySchema, AgentMethodSchema, AgentTypeSchema, AutoInjectedKind, FieldSource, InputSchema, NamedField, OutputSchema, - ParsedAgentId, json_input_schema_value_to_typed_schema_value, + ParsedAgentId, contains_stream_in_graph, json_input_schema_value_to_typed_schema_value, 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 proptest::prelude::*; use serde_json::json; use test_r::test; @@ -652,7 +653,7 @@ fn projected_helper_drops_all_defs_for_ref_free_root() { assert!(typed.graph().defs.is_empty()); } -// --- AgentTypeSchema::validate: rejection of P3 stub types (future/stream) --- +// --- AgentTypeSchema::validate: stream placement and shared definitions --- fn method(name: &str, input: Vec, output: OutputSchema) -> AgentMethodSchema { AgentMethodSchema { @@ -680,22 +681,109 @@ fn validate_accepts_agent_type_without_p3_stubs() { } #[test] -fn validate_rejects_stream_in_method_output() { +fn validate_accepts_streams_in_method_input_and_output() { let agent = AgentTypeSchema { methods: vec![method( - "download", - vec![], + "transform", + vec![NamedField::user_supplied( + "input", + SchemaType::stream(Some(SchemaType::string())), + )], OutputSchema::Single(Box::new(SchemaType::stream(Some(SchemaType::u8())))), )], ..sample_agent_type() }; - let err = agent - .validate() - .expect_err("stream output must be rejected"); + agent.validate().expect("method streams must be accepted"); +} + +#[test] +fn stream_classification_ignores_definitions_unreachable_from_the_method() { + let graph = registry(vec![ + proj_def("Plain", SchemaType::string()), + proj_def( + "UnrelatedStream", + SchemaType::stream(Some(SchemaType::u8())), + ), + ]); + let method = method( + "plain", + vec![NamedField::user_supplied( + "input", + SchemaType::ref_to(TypeId::new("Plain")), + )], + OutputSchema::Single(Box::new(SchemaType::u64())), + ); + assert!( - err.contains("unsupported type 'stream'") && err.contains("method 'download' output"), - "unexpected error: {err}" + method + .input_schema + .fields() + .iter() + .all(|field| !contains_stream_in_graph(&graph, &field.schema)) ); + assert!( + !contains_stream_in_graph(&graph, method.output_schema.schema().unwrap()), + "the unrelated stream definition must not classify the method as streaming" + ); + assert!(contains_stream_in_graph( + &graph, + &SchemaType::ref_to(TypeId::new("UnrelatedStream")) + )); + assert!(!method.uses_streams(&graph)); +} + +#[test] +fn method_stream_classification_follows_input_and_output_refs() { + let graph = registry(vec![ + proj_def( + "Input", + SchemaType::record(vec![proj_field( + "items", + SchemaType::stream(Some(SchemaType::string())), + )]), + ), + proj_def( + "Output", + SchemaType::option(SchemaType::stream(Some(SchemaType::u8()))), + ), + ]); + let input_method = method( + "input", + vec![NamedField::user_supplied( + "input", + SchemaType::ref_to(TypeId::new("Input")), + )], + OutputSchema::Unit, + ); + let output_method = method( + "output", + vec![], + OutputSchema::Single(Box::new(SchemaType::ref_to(TypeId::new("Output")))), + ); + + assert!(input_method.uses_streams(&graph)); + assert!(output_method.uses_streams(&graph)); +} + +#[test] +fn method_input_validation_accepts_a_live_stream_handle() { + let method = method( + "consume", + vec![NamedField::user_supplied( + "input", + SchemaType::stream(Some(SchemaType::u32())), + )], + OutputSchema::Unit, + ); + let input = SchemaValue::Record { + fields: vec![SchemaValue::Stream(SchemaValueStream::from_host_endpoint( + (), + ))], + }; + + method + .validate_input(&SchemaGraph::empty(), &input) + .unwrap(); } #[test] @@ -715,8 +803,50 @@ fn validate_rejects_future_in_method_parameter() { .validate() .expect_err("future parameter must be rejected"); assert!( - err.contains("unsupported type 'future'") - && err.contains("method 'wait-for' parameter 'signal'"), + err.contains("future values are not allowed in scope AgentMethodInput"), + "unexpected error: {err}" + ); +} + +#[test] +fn validate_rejects_future_in_unreferenced_shared_definition() { + let agent = AgentTypeSchema { + schema: registry(vec![proj_def( + "UnusedFuture", + SchemaType::future(Some(SchemaType::string())), + )]), + ..sample_agent_type() + }; + let err = agent + .validate() + .expect_err("future must be rejected even in an unreferenced definition"); + assert!( + err.contains( + "shared type definition 'UnusedFuture' contains the unsupported type 'future'" + ), + "unexpected error: {err}" + ); +} + +#[test] +fn validate_rejects_stream_in_auto_injected_method_parameter() { + let agent = AgentTypeSchema { + methods: vec![method( + "invalid", + vec![NamedField::auto_injected( + "principal", + AutoInjectedKind::Principal, + SchemaType::stream(Some(SchemaType::string())), + )], + OutputSchema::Unit, + )], + ..sample_agent_type() + }; + let err = agent + .validate() + .expect_err("auto-injected streams must be rejected"); + assert!( + err.contains("stream values are not allowed in scope Boundary"), "unexpected error: {err}" ); } @@ -740,29 +870,43 @@ fn validate_rejects_stream_nested_in_constructor_parameter() { .validate() .expect_err("nested stream in constructor must be rejected"); assert!( - err.contains("unsupported type 'stream'") - && err.contains("constructor parameter 'sources'"), + err.contains("stream values are not allowed in scope Constructor"), "unexpected error: {err}" ); } #[test] -fn validate_rejects_stream_in_shared_type_definition() { +fn validate_accepts_stream_in_shared_type_definition_used_by_method() { let agent = AgentTypeSchema { schema: registry(vec![proj_def( "Chunks", SchemaType::stream(Some(SchemaType::u8())), )]), + methods: vec![method( + "download", + vec![], + OutputSchema::Single(Box::new(SchemaType::ref_to(TypeId::new("Chunks")))), + )], ..sample_agent_type() }; - let err = agent + agent .validate() - .expect_err("stream in a shared type definition must be rejected"); - assert!( - err.contains("unsupported type 'stream'") - && err.contains("shared type definition 'Chunks'"), - "unexpected error: {err}" - ); + .expect("method may reference a shared stream definition"); +} + +#[test] +fn validate_rejects_stream_in_unreferenced_shared_definition() { + let agent = AgentTypeSchema { + schema: registry(vec![proj_def( + "UnusedStream", + SchemaType::stream(Some(SchemaType::string())), + )]), + ..sample_agent_type() + }; + + agent + .validate() + .expect_err("streams are allowed only at explicit agent method input/output positions"); } #[test] @@ -779,13 +923,13 @@ fn validate_rejects_future_in_config_value_type() { .validate() .expect_err("future config value type must be rejected"); assert!( - err.contains("unsupported type 'future'") && err.contains("config value at path 'a.b'"), + err.contains("future values are not allowed in scope Boundary"), "unexpected error: {err}" ); } #[test] -fn validate_rejects_stream_in_dependency_method_parameter() { +fn validate_accepts_stream_in_dependency_method_parameter() { let agent = AgentTypeSchema { dependencies: vec![AgentDependencySchema { type_name: "helper".to_string(), @@ -808,12 +952,7 @@ fn validate_rejects_stream_in_dependency_method_parameter() { }], ..sample_agent_type() }; - let err = agent + agent .validate() - .expect_err("stream in a dependency method must be rejected"); - assert!( - err.contains("unsupported type 'stream'") - && err.contains("dependency 'helper' method 'feed' parameter 'data'"), - "unexpected error: {err}" - ); + .expect("dependency method streams must be accepted"); } diff --git a/golem-common/src/schema/mod.rs b/golem-common/src/schema/mod.rs index 55ddef7cc5..7757b2f934 100644 --- a/golem-common/src/schema/mod.rs +++ b/golem-common/src/schema/mod.rs @@ -33,7 +33,7 @@ pub use golem_schema::schema::proptest_strategies; pub use golem_schema::schema::wit; pub use golem_schema::schema::{ canonical, conversion, derive, graph, host_managed, metadata, multimodal, schema_type, - schema_value, unstructured, + schema_value, stream, unstructured, }; #[cfg(test)] diff --git a/golem-common/src/schema/render/json_value.rs b/golem-common/src/schema/render/json_value.rs index 33fdc8405d..75591958e9 100644 --- a/golem-common/src/schema/render/json_value.rs +++ b/golem-common/src/schema/render/json_value.rs @@ -1095,6 +1095,7 @@ fn value_name(value: &SchemaValue) -> &'static str { SchemaValue::Union(_) => "union", SchemaValue::Secret(_) => "secret", SchemaValue::QuotaToken(_) => "quota-token", + SchemaValue::Stream(_) => "stream", } } diff --git a/golem-common/src/schema/tests/protobuf_tests.rs b/golem-common/src/schema/tests/protobuf_tests.rs index 42c8e2652b..4d72f6a962 100644 --- a/golem-common/src/schema/tests/protobuf_tests.rs +++ b/golem-common/src/schema/tests/protobuf_tests.rs @@ -57,7 +57,8 @@ proptest! { /// quantity / secret / quota-token). #[test] fn schema_value_proto_round_trip(value in schema_value_strategy()) { - let proto: golem_api_grpc::proto::golem::schema::SchemaValue = value.clone().into(); + let proto: golem_api_grpc::proto::golem::schema::SchemaValue = + value.clone().try_into().expect("encode"); let back: SchemaValue = proto.try_into().expect("decode"); prop_assert!( schema_values_eq(&value, &back), @@ -68,7 +69,8 @@ proptest! { /// The typed pair (graph + value) round-trips through its protobuf mirror. #[test] fn typed_schema_value_proto_round_trip(typed in typed_schema_value_strategy()) { - let proto: golem_api_grpc::proto::golem::schema::TypedSchemaValue = typed.clone().into(); + let proto: golem_api_grpc::proto::golem::schema::TypedSchemaValue = + typed.clone().try_into().expect("encode"); let back: TypedSchemaValue = proto.try_into().expect("decode"); prop_assert_eq!(typed.graph(), back.graph()); prop_assert!( diff --git a/golem-common/src/schema/tests/wit_tests.rs b/golem-common/src/schema/tests/wit_tests.rs index a76cd1d766..470e5be19c 100644 --- a/golem-common/src/schema/tests/wit_tests.rs +++ b/golem-common/src/schema/tests/wit_tests.rs @@ -18,13 +18,14 @@ use crate::schema::proptest_strategies as strategies; use crate::schema::schema_type::SchemaType; use crate::schema::schema_value::{QuotaTokenValuePayload, SchemaValue, SecretValuePayload}; use crate::schema::wit::{ - DecodeError, EncodeError, QuotaTokenHandleRep, QuotaTokenResolver, SecretHandleRep, - SecretResolver, decode_graph, decode_typed, decode_typed_rejecting_quota_with, decode_value, - decode_value_rejecting_quota_with, decode_value_with, encode_graph, encode_typed, encode_value, - encode_value_with, wire, + DecodeError, EncodeError, QuotaTokenHandleRep, QuotaTokenResolver, SchemaValueStreamResolver, + SecretHandleRep, SecretResolver, decode_graph, decode_typed, decode_typed_rejecting_quota_with, + decode_value, decode_value_rejecting_quota_with, decode_value_with, encode_graph, encode_typed, + encode_value, encode_value_with, wire, }; use chrono::{TimeZone, Utc}; use golem_schema::model::EnvironmentId; +use golem_schema::schema::{SchemaValueStream, SchemaValueStreamHandleRep}; use proptest::prelude::*; use strategies::{ schema_graph_strategy, transportable_schema_value_strategy, @@ -381,6 +382,34 @@ impl SecretResolver for TableResolver { } } +impl SchemaValueStreamResolver for TableResolver { + type Error = anyhow::Error; + + fn handle_from_stream( + &mut self, + stream: SchemaValueStream, + ) -> Result, Self::Error> { + let handle = self.table.push(SchemaValueStreamHandleRep::new(stream))?; + self.live += 1; + Ok(handle) + } + + fn stream_from_handle( + &mut self, + handle: Resource, + ) -> Result { + let stream = self.table.delete(handle)?.into_stream(); + self.live -= 1; + Ok(stream) + } + + fn drop_stream_handle(&mut self, handle: Resource) { + if self.table.delete(handle).is_ok() { + self.live -= 1; + } + } +} + #[test] fn quota_token_round_trips_through_resolver() { let value = SchemaValue::QuotaToken(sample_snapshot()); @@ -402,6 +431,25 @@ fn secret_round_trips_through_resolver() { assert_eq!(resolver.live, 0); } +#[test] +fn stream_decodes_through_the_standard_resolver_path() { + let stream = SchemaValueStream::from_host_endpoint(42_u32); + let mut resolver = TableResolver::new(); + let handle = resolver.handle_from_stream(stream).unwrap(); + let tree = wire::SchemaValueTree { + value_nodes: vec![wire::SchemaValueNode::StreamValue(handle)], + root: 0, + }; + + let decoded = decode_value_with(tree, &mut resolver).expect("decode stream"); + let SchemaValue::Stream(decoded) = decoded else { + panic!("expected stream value"); + }; + + assert_eq!(decoded.take_host_endpoint::().unwrap(), 42); + assert_eq!(resolver.live, 0); +} + #[test] fn nested_quota_token_round_trips_through_resolver() { let value = SchemaValue::Record { diff --git a/golem-common/src/schema/validation/placement.rs b/golem-common/src/schema/validation/placement.rs index 39403c4845..fa439094e4 100644 --- a/golem-common/src/schema/validation/placement.rs +++ b/golem-common/src/schema/validation/placement.rs @@ -24,15 +24,15 @@ //! (typed values, `Custom` oplog payloads, REST/RPC envelopes). //! - [`validate_agent_type_placement`] walks an [`AgentTypeSchema`]: the //! constructor input fields are validated against -//! [`SchemaScope::Constructor`], method inputs and outputs against -//! [`SchemaScope::Boundary`], the agent's named defs against -//! [`SchemaScope::Boundary`], and each dependency recursively against -//! its own graph. The sentinel `graph.root` on agent carriers is not -//! walked. +//! [`SchemaScope::Constructor`], method inputs and outputs against their +//! stream-capable method scopes, and each dependency recursively against its +//! own graph. Named definitions are validated at each use site because the +//! same definition may be legal in a method and illegal in a constructor. +//! The sentinel `graph.root` on agent carriers is not walked. use crate::schema::agent::{ - AgentConstructorSchema, AgentDependencySchema, AgentMethodSchema, AgentTypeSchema, InputSchema, - OutputSchema, + AgentConstructorSchema, AgentDependencySchema, AgentMethodSchema, AgentTypeSchema, FieldSource, + InputSchema, OutputSchema, }; use crate::schema::graph::SchemaGraph; use crate::schema::host_managed::HostManagedKind; @@ -51,6 +51,10 @@ pub enum SchemaScope { Persisted, /// REST / RPC boundary payloads. Boundary, + /// Agent method inputs, whose recursive value trees may contain streams. + AgentMethodInput, + /// Agent method outputs, whose recursive value trees may contain streams. + AgentMethodOutput, /// Public docs / schema rendering. Docs, /// User-provided `Custom` durable payloads. @@ -69,6 +73,11 @@ pub enum PlacementError { /// A field / definition annotated with [`Role::Multimodal`] whose body /// is `list>` appeared in [`SchemaScope::Constructor`]. MultimodalListNotAllowedInConstructor, + /// A [`SchemaType::Stream`] node appeared outside a stream-capable method. + StreamNotAllowed { scope: SchemaScope }, + /// A [`SchemaType::Future`] node appeared in a schema boundary. Future + /// values are not part of the agent ABI. + FutureNotAllowed { scope: SchemaScope }, } impl Display for PlacementError { @@ -84,6 +93,12 @@ impl Display for PlacementError { f, "a multimodal `list>` is not allowed in constructor scope" ), + PlacementError::StreamNotAllowed { scope } => { + write!(f, "stream values are not allowed in scope {scope:?}") + } + PlacementError::FutureNotAllowed { scope } => { + write!(f, "future values are not allowed in scope {scope:?}") + } } } } @@ -144,6 +159,21 @@ fn walk_type<'a>( errors: &mut Vec, visited: &mut Vec<&'a TypeId>, ) { + match ty { + SchemaType::Stream { .. } + if !matches!( + scope, + SchemaScope::AgentMethodInput | SchemaScope::AgentMethodOutput + ) => + { + errors.push(PlacementError::StreamNotAllowed { scope }); + } + SchemaType::Future { .. } => { + errors.push(PlacementError::FutureNotAllowed { scope }); + } + _ => {} + } + // Constructor-scope check: a list> tagged anywhere on its // metadata-carrying nodes (enclosing field/def metadata, list node // metadata, inner element Ref metadata, or inner variant metadata) with @@ -376,14 +406,14 @@ fn resolve_ref_chain<'a>( /// /// Walks: /// - the constructor input fields against [`SchemaScope::Constructor`] -/// - each method's input and output bodies against [`SchemaScope::Boundary`] -/// - every named def in [`AgentTypeSchema::schema`] against -/// [`SchemaScope::Boundary`] +/// - each method's input and output bodies against the corresponding +/// stream-capable method scope +/// - config value types against the scalar-only [`SchemaScope::Boundary`] /// - each dependency, recursively, against its own /// [`AgentDependencySchema::schema`] /// -/// The sentinel `graph.root` carried by agent-layer [`SchemaGraph`]s is -/// **not** walked (see §4.22). +/// Named definitions are resolved and checked at each use site. The sentinel +/// `graph.root` carried by agent-layer [`SchemaGraph`]s is **not** walked. /// /// This is a **placement-only** validator. Structural well-formedness /// (dangling refs, duplicate ids, …) is the responsibility of @@ -396,7 +426,18 @@ pub fn validate_agent_type_placement(ty: &AgentTypeSchema) -> Result<(), Vec, ) { - walk_input_schema(graph, &method.input_schema, SchemaScope::Boundary, errors); - walk_output_schema(graph, &method.output_schema, SchemaScope::Boundary, errors); + match &method.input_schema { + InputSchema::Parameters(fields) => { + for field in fields { + let scope = match field.source { + FieldSource::UserSupplied => SchemaScope::AgentMethodInput, + FieldSource::AutoInjected(_) => SchemaScope::Boundary, + }; + let mut visited = Vec::new(); + walk_type( + graph, + &field.schema, + &field.metadata, + scope, + errors, + &mut visited, + ); + } + } + } + walk_output_schema( + graph, + &method.output_schema, + SchemaScope::AgentMethodOutput, + errors, + ); } fn walk_input_schema( @@ -483,26 +546,3 @@ fn walk_output_schema( } } } - -/// Walk every def in `graph.defs` at [`SchemaScope::Boundary`]. -/// -/// Defs are validated at Boundary because shared defs may legitimately be -/// used by method inputs/outputs, where Constructor-only restrictions -/// (e.g. `Secret`) do not apply. Constructor-specific restrictions are -/// still enforced at constructor use sites: when a constructor parameter -/// is a [`SchemaType::Ref`], the constructor walk resolves the ref and -/// re-checks the resolved body under [`SchemaScope::Constructor`]. -fn walk_agent_graph_defs(graph: &SchemaGraph, errors: &mut Vec) { - for def in &graph.defs { - let body_metadata = def.body.metadata().clone(); - let mut visited: Vec<&TypeId> = Vec::new(); - walk_type( - graph, - &def.body, - &body_metadata, - SchemaScope::Boundary, - errors, - &mut visited, - ); - } -} diff --git a/golem-common/src/schema/validation/tests/placement_tests.rs b/golem-common/src/schema/validation/tests/placement_tests.rs index ac178b0721..1803eca2f9 100644 --- a/golem-common/src/schema/validation/tests/placement_tests.rs +++ b/golem-common/src/schema/validation/tests/placement_tests.rs @@ -157,6 +157,8 @@ fn plain_primitives_allowed_in_every_scope() { SchemaScope::Constructor, SchemaScope::Persisted, SchemaScope::Boundary, + SchemaScope::AgentMethodInput, + SchemaScope::AgentMethodOutput, SchemaScope::Docs, SchemaScope::Custom, ] { @@ -164,6 +166,44 @@ fn plain_primitives_allowed_in_every_scope() { } } +#[test] +fn streams_are_allowed_only_on_agent_method_boundaries() { + let graph = SchemaGraph::anonymous(SchemaType::stream(Some(SchemaType::string()))); + for scope in [ + SchemaScope::Constructor, + SchemaScope::Persisted, + SchemaScope::Boundary, + SchemaScope::Docs, + SchemaScope::Custom, + ] { + let errors = validate_placement(&graph, scope).expect_err("stream must be rejected"); + assert!(errors.contains(&PlacementError::StreamNotAllowed { scope })); + } + for scope in [ + SchemaScope::AgentMethodInput, + SchemaScope::AgentMethodOutput, + ] { + validate_placement(&graph, scope).expect("stream must be accepted"); + } +} + +#[test] +fn futures_are_rejected_in_every_scope() { + let graph = SchemaGraph::anonymous(SchemaType::future(Some(SchemaType::string()))); + for scope in [ + SchemaScope::Constructor, + SchemaScope::Persisted, + SchemaScope::Boundary, + SchemaScope::AgentMethodInput, + SchemaScope::AgentMethodOutput, + SchemaScope::Docs, + SchemaScope::Custom, + ] { + let errors = validate_placement(&graph, scope).expect_err("future must be rejected"); + assert!(errors.contains(&PlacementError::FutureNotAllowed { scope })); + } +} + // -------------------------------------------------------------------------- // Agent-aware placement // -------------------------------------------------------------------------- @@ -174,7 +214,7 @@ mod agent { use crate::base_model::agent::{AgentMode, AgentTypeName, Snapshotting}; use crate::schema::agent::{ AgentConstructorSchema, AgentDependencySchema, AgentMethodSchema, AgentTypeSchema, - InputSchema, NamedField, OutputSchema, + AutoInjectedKind, InputSchema, NamedField, OutputSchema, }; use crate::schema::schema_type::SecretSpec; use crate::schema::validation::placement::{ @@ -209,6 +249,27 @@ mod agent { validate_agent_type_placement(&agent).expect("empty agent should pass"); } + #[test] + fn auto_injected_method_stream_fails_agent_placement_validation() { + let mut agent = empty_agent("a"); + agent.methods.push(AgentMethodSchema { + name: "invalid".into(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![NamedField::auto_injected( + "principal", + AutoInjectedKind::Principal, + SchemaType::stream(Some(SchemaType::string())), + )]), + output_schema: OutputSchema::Unit, + http_endpoint: vec![], + read_only: None, + }); + + validate_agent_type_placement(&agent) + .expect_err("method input streams must be in caller-supplied fields"); + } + #[test] fn secret_in_constructor_input_is_rejected_at_agent_layer() { let mut agent = empty_agent("a"); diff --git a/golem-common/wit/deps/golem-agent/guest.wit b/golem-common/wit/deps/golem-agent/guest.wit index dbe97ff310..d72413c288 100644 --- a/golem-common/wit/deps/golem-agent/guest.wit +++ b/golem-common/wit/deps/golem-agent/guest.wit @@ -15,8 +15,9 @@ interface guest { /// Invokes an agent. If create was not called before, it fails. /// /// `input` is a value tree whose root encodes the method's parameter list. - /// The result is `none` when the method's `output-schema` is `unit`, and - /// `some(value)` for a `single` output. + /// Streams are represented recursively by `stream-value` nodes. The result + /// is `none` when the method's `output-schema` is `unit`, and `some(value)` + /// for a `single` output. invoke: async func(method-name: string, input: schema-value-tree, principal: principal) -> result, agent-error>; /// Gets the agent type. If create was not called before, it fails diff --git a/golem-common/wit/deps/golem-core-v2/golem-core-v2.wit b/golem-common/wit/deps/golem-core-v2/golem-core-v2.wit index f393d304d6..2b731300f4 100644 --- a/golem-common/wit/deps/golem-core-v2/golem-core-v2.wit +++ b/golem-common/wit/deps/golem-core-v2/golem-core-v2.wit @@ -125,6 +125,17 @@ interface types { /// and reveal it only through capability-gated host interfaces. resource secret; + /// An affine wrapper around a native Component Model stream of schema + /// values. The indirection lets a stream occur anywhere in a recursive + /// `schema-value-tree` while preserving the native stream endpoint. + resource schema-value-stream { + /// Wraps any native schema-value reader for placement in a value tree. + wrap: static async func(reader: stream) -> own; + + /// Consumes a schema-value-stream wrapper and returns its native reader. + unwrap: static async func(value: own) -> stream; + } + // ============================================================ // Schema graph (self-contained type carrier) // ============================================================ @@ -549,8 +560,7 @@ interface types { // Capability nodes secret-value(own), quota-token-handle(own), - - // WASI P3 stubs (parseable only in the schema; no constructible values). + stream-value(own), } record variant-value-payload { @@ -609,4 +619,5 @@ interface types { graph: schema-graph, value: schema-value-tree, } + } diff --git a/golem-debugging-service/config/debug-worker-executor.sample.env b/golem-debugging-service/config/debug-worker-executor.sample.env index 6755e9bbe0..2acb396b8b 100644 --- a/golem-debugging-service/config/debug-worker-executor.sample.env +++ b/golem-debugging-service/config/debug-worker-executor.sample.env @@ -61,6 +61,7 @@ GOLEM__LIMITS__EVENT_BROADCAST_CAPACITY=1024 GOLEM__LIMITS__EVENT_HISTORY_SIZE=128 GOLEM__LIMITS__FUEL_TO_BORROW=10000 GOLEM__LIMITS__INVOCATION_RESULT_BROADCAST_CAPACITY=100000 +GOLEM__LIMITS__LIVE_STREAM_EVENT_BROADCAST_CAPACITY=32 GOLEM__LIMITS__MAX_ACTIVE_WORKERS=1024 GOLEM__LIMITS__MAX_CONCURRENT_STREAMS=1024 GOLEM__LIMITS__MAX_INVOCATION_CONTEXT_STACK_DEPTH=1024 @@ -251,6 +252,7 @@ GOLEM__LIMITS__EVENT_BROADCAST_CAPACITY=1024 GOLEM__LIMITS__EVENT_HISTORY_SIZE=128 GOLEM__LIMITS__FUEL_TO_BORROW=10000 GOLEM__LIMITS__INVOCATION_RESULT_BROADCAST_CAPACITY=100000 +GOLEM__LIMITS__LIVE_STREAM_EVENT_BROADCAST_CAPACITY=32 GOLEM__LIMITS__MAX_ACTIVE_WORKERS=1024 GOLEM__LIMITS__MAX_CONCURRENT_STREAMS=1024 GOLEM__LIMITS__MAX_INVOCATION_CONTEXT_STACK_DEPTH=1024 diff --git a/golem-debugging-service/config/debug-worker-executor.toml b/golem-debugging-service/config/debug-worker-executor.toml index e1109502e8..5b778e83d0 100644 --- a/golem-debugging-service/config/debug-worker-executor.toml +++ b/golem-debugging-service/config/debug-worker-executor.toml @@ -104,6 +104,7 @@ event_broadcast_capacity = 1024 event_history_size = 128 fuel_to_borrow = 10000 invocation_result_broadcast_capacity = 100000 +live_stream_event_broadcast_capacity = 32 max_active_workers = 1024 max_concurrent_streams = 1024 max_invocation_context_stack_depth = 1024 @@ -392,6 +393,7 @@ without_time = false # event_history_size = 128 # fuel_to_borrow = 10000 # invocation_result_broadcast_capacity = 100000 +# live_stream_event_broadcast_capacity = 32 # max_active_workers = 1024 # max_concurrent_streams = 1024 # max_invocation_context_stack_depth = 1024 diff --git a/golem-debugging-service/src/lib.rs b/golem-debugging-service/src/lib.rs index ef2c848dc9..d4e03d81af 100644 --- a/golem-debugging-service/src/lib.rs +++ b/golem-debugging-service/src/lib.rs @@ -33,7 +33,7 @@ use async_trait::async_trait; use golem_service_base::clients::registry::RegistryService; use golem_service_base::storage::blob::BlobStorage; pub use golem_worker_executor::RunDetails; -use golem_worker_executor::durable_host::DurableWorkerCtx; +use golem_worker_executor::durable_host::{CoreTypesHost, DurableWorkerCtx}; use golem_worker_executor::preview2::{golem_api_1_x, golem_durability}; use golem_worker_executor::services::active_workers::ActiveWorkers; use golem_worker_executor::services::agent_types::AgentTypesService; @@ -251,9 +251,13 @@ pub async fn create_debugging_service_services( // When it comes to fork, we need the original oplog service let worker_fork = Arc::new(DefaultWorkerFork::new( - Arc::new(RemoteInvocationRpc::new( + Arc::new(RemoteInvocationRpc::new_with_stream_capacity( worker_proxy.clone(), shard_service.clone(), + golem_config + .limits + .live_stream_event_broadcast_capacity + .get(), )), active_workers.clone(), engine.clone(), @@ -293,9 +297,13 @@ pub async fn create_debugging_service_services( )); let rpc = Arc::new(DirectWorkerInvocationRpc::new( - Arc::new(RemoteInvocationRpc::new( + Arc::new(RemoteInvocationRpc::new_with_stream_capacity( worker_proxy.clone(), shard_service.clone(), + golem_config + .limits + .live_stream_event_broadcast_capacity + .get(), )), direct_invocation_auth_service, active_workers.clone(), @@ -487,7 +495,7 @@ pub fn create_debug_wasmtime_linker(engine: &Engine) -> anyhow::Result>, >(&mut linker, get_durable_ctx)?; - golem_schema::schema::wit::wire::add_to_linker::<_, HasSelf>>( + golem_schema::schema::wit::wire::add_to_linker::<_, CoreTypesHost>( &mut linker, get_durable_ctx, )?; diff --git a/golem-openapi-client-generator/src/rust/client_gen.rs b/golem-openapi-client-generator/src/rust/client_gen.rs index 8a2928de08..242a1c6dfd 100644 --- a/golem-openapi-client-generator/src/rust/client_gen.rs +++ b/golem-openapi-client-generator/src/rust/client_gen.rs @@ -260,7 +260,9 @@ fn tag_operations( if let Some(item) = path_item.as_item() { if let Some(tag) = tag { item.iter() - .filter(|(_, op)| op.tags.contains(&tag.name)) + .filter(|(_, op)| { + op.tags.contains(&tag.name) && !is_websocket_upgrade_operation(op) + }) .map(|(method, op)| PathOperation { path: Path::from_string(path), original_path: path.to_string(), @@ -270,7 +272,7 @@ fn tag_operations( .collect() } else { item.iter() - .filter(|(_, op)| op.tags.is_empty()) + .filter(|(_, op)| op.tags.is_empty() && !is_websocket_upgrade_operation(op)) .map(|(method, op)| PathOperation { path: Path::from_string(path), original_path: path.to_string(), @@ -284,6 +286,13 @@ fn tag_operations( } } +fn is_websocket_upgrade_operation(operation: &Operation) -> bool { + operation + .responses + .responses + .contains_key(&StatusCode::Code(101)) +} + fn match_tag(tag: &Option, path_item: &ReferenceOr) -> bool { if let Some(item) = path_item.as_item() { if let Some(tag) = tag { diff --git a/golem-registry-service/src/grpc/api_impl.rs b/golem-registry-service/src/grpc/api_impl.rs index 83ce75825f..7a6d8fa001 100644 --- a/golem-registry-service/src/grpc/api_impl.rs +++ b/golem-registry-service/src/grpc/api_impl.rs @@ -295,7 +295,7 @@ impl RegistryServiceGrpcApi { .await?; Ok(GetComponentMetadataSuccessResponse { - component: Some(component.into()), + component: Some(component.try_into()?), }) } @@ -314,7 +314,7 @@ impl RegistryServiceGrpcApi { .await?; Ok(GetDeployedComponentMetadataSuccessResponse { - component: Some(component.into()), + component: Some(component.try_into()?), }) } @@ -333,7 +333,10 @@ impl RegistryServiceGrpcApi { .await?; Ok(GetAllDeployedComponentRevisionsSuccessResponse { - components: components.into_iter().map(|c| c.into()).collect(), + components: components + .into_iter() + .map(TryInto::try_into) + .collect::>()?, }) } @@ -370,7 +373,7 @@ impl RegistryServiceGrpcApi { .await?; Ok(ResolveComponentSuccessResponse { - component: Some(component.into()), + component: Some(component.try_into()?), }) } @@ -517,7 +520,7 @@ impl RegistryServiceGrpcApi { .await?; Ok(GetCurrentEnvironmentStateSuccessResponse { - environment_state: Some(environment_state.into()), + environment_state: Some(environment_state.try_into()?), }) } @@ -554,7 +557,7 @@ impl RegistryServiceGrpcApi { retry_policies: Vec::new(), tool_deployment: None, } - .into(), + .try_into()?, ), }) } diff --git a/golem-schema-derive/src/codegen/poem.rs b/golem-schema-derive/src/codegen/poem.rs index 3c8c350b3c..eb45de5b81 100644 --- a/golem-schema-derive/src/codegen/poem.rs +++ b/golem-schema-derive/src/codegen/poem.rs @@ -59,6 +59,7 @@ struct SerdeFieldAttrs { #[derive(Default)] struct SerdeVariantAttrs { rename: Option, + skip: bool, } pub fn expand_poem_schema(input: &DeriveInput) -> syn::Result { @@ -222,6 +223,9 @@ fn expand_adjacent_enum( let mut variant_blocks = Vec::new(); for variant in &data.variants { let vattrs = parse_serde_variant_attrs(&variant.attrs)?; + if vattrs.skip { + continue; + } let case = resolve_variant_name(&variant.ident, &vattrs, serde.rename_all); let case_lit = LitStr::new(&case, variant.ident.span()); @@ -983,6 +987,7 @@ fn parse_serde_variant_attrs(attrs: &[Attribute]) -> syn::Result()?.value()); } + "skip" => out.skip = true, other => { return Err(meta.error(format!( "PoemSchema does not support `#[serde({other})]` on a variant" diff --git a/golem-schema/Cargo.toml b/golem-schema/Cargo.toml index edc7e8359a..415fd65e86 100644 --- a/golem-schema/Cargo.toml +++ b/golem-schema/Cargo.toml @@ -36,6 +36,7 @@ golem-schema-derive = { workspace = true } base64 = { workspace = true } bigdecimal = { workspace = true } bit-vec = { workspace = true } +blake3 = { workspace = true } bytes = { workspace = true, optional = true } chrono = { workspace = true } combine = { workspace = true } diff --git a/golem-schema/src/schema/conversion.rs b/golem-schema/src/schema/conversion.rs index 5ff1b7c337..76043db3fb 100644 --- a/golem-schema/src/schema/conversion.rs +++ b/golem-schema/src/schema/conversion.rs @@ -460,6 +460,7 @@ pub fn value_kind(v: &SchemaValue) -> &'static str { SchemaValue::Union(_) => "union", SchemaValue::Secret(_) => "secret", SchemaValue::QuotaToken(_) => "quota-token", + SchemaValue::Stream(_) => "stream", } } diff --git a/golem-schema/src/schema/mod.rs b/golem-schema/src/schema/mod.rs index b57421ad0f..67158daa55 100644 --- a/golem-schema/src/schema/mod.rs +++ b/golem-schema/src/schema/mod.rs @@ -25,6 +25,7 @@ pub mod multimodal; pub mod protobuf; pub mod schema_type; pub mod schema_value; +pub mod stream; pub mod tool; pub mod unstructured; pub mod validation; @@ -56,3 +57,6 @@ pub use schema_value::{ BinaryValuePayload, DurationValuePayload, QuotaTokenValuePayload, ResultValuePayload, SchemaValue, SecretValuePayload, TextValuePayload, UnionValuePayload, VariantValuePayload, }; +pub use stream::SchemaValueStream; +#[cfg(all(feature = "host", not(feature = "guest")))] +pub use stream::SchemaValueStreamHandleRep; diff --git a/golem-schema/src/schema/protobuf.rs b/golem-schema/src/schema/protobuf.rs index 4eafeaa3cd..5b46e04274 100644 --- a/golem-schema/src/schema/protobuf.rs +++ b/golem-schema/src/schema/protobuf.rs @@ -15,6 +15,7 @@ //! Conversions between the recursive in-memory schema model and its protobuf //! mirror in the `golem.schema` package. +#[cfg(not(all(feature = "guest", not(feature = "host"))))] use crate::model::EnvironmentId; use crate::schema::graph::{SchemaGraph, SchemaTypeDef, TypedSchemaValue}; use crate::schema::metadata::{MetadataEnvelope, Role, TypeId}; @@ -25,9 +26,11 @@ use crate::schema::schema_type::{ UrlRestrictions, VariantCaseType, }; use crate::schema::schema_value::{ - BinaryValuePayload, DurationValuePayload, QuotaTokenValuePayload, ResultValuePayload, - SchemaValue, SecretValuePayload, TextValuePayload, UnionValuePayload, VariantValuePayload, + BinaryValuePayload, DurationValuePayload, ResultValuePayload, SchemaValue, TextValuePayload, + UnionValuePayload, VariantValuePayload, }; +#[cfg(not(all(feature = "guest", not(feature = "host"))))] +use crate::schema::schema_value::{QuotaTokenValuePayload, SecretValuePayload}; use chrono::{DateTime, TimeZone, Utc}; use golem_api_grpc::proto::golem::common::Empty as ProtoEmpty; use golem_api_grpc::proto::golem::schema as proto; @@ -87,8 +90,8 @@ fn req_box_from_proto( // --- value-side helpers ------------------------------------------------------ -fn value_to_boxed_proto(value: SchemaValue) -> Box { - Box::new(value.into()) +fn value_to_boxed_proto(value: SchemaValue) -> Result, String> { + Ok(Box::new(value.try_into()?)) } fn opt_box_value_from_proto( @@ -908,8 +911,10 @@ impl TryFrom for SchemaTypeDef { // --- SchemaValue / TypedSchemaValue ------------------------------------------ -impl From for proto::SchemaValue { - fn from(value: SchemaValue) -> Self { +impl TryFrom for proto::SchemaValue { + type Error = String; + + fn try_from(value: SchemaValue) -> Result { let body = match value { SchemaValue::Bool(b) => ValueBody::BoolValue(b), SchemaValue::S8(v) => ValueBody::S8Value(v as i32), @@ -925,45 +930,64 @@ impl From for proto::SchemaValue { SchemaValue::Char(c) => ValueBody::CharValue(c as u32), SchemaValue::String(s) => ValueBody::StringValue(s), SchemaValue::Record { fields } => ValueBody::RecordValue(proto::RecordValue { - fields: fields.into_iter().map(Into::into).collect(), + fields: fields + .into_iter() + .map(TryInto::try_into) + .collect::>()?, }), SchemaValue::Variant(p) => ValueBody::VariantValue(Box::new(proto::VariantValue { case: p.case, - payload: p.payload.map(|value| value_to_boxed_proto(*value)), + payload: p + .payload + .map(|value| value_to_boxed_proto(*value)) + .transpose()?, })), SchemaValue::Enum { case } => ValueBody::EnumValue(case), SchemaValue::Flags { bits } => ValueBody::FlagsValue(proto::FlagsValue { bits }), SchemaValue::Tuple { elements } => ValueBody::TupleValue(proto::TupleValue { - elements: elements.into_iter().map(Into::into).collect(), + elements: elements + .into_iter() + .map(TryInto::try_into) + .collect::>()?, }), SchemaValue::List { elements } => ValueBody::ListValue(proto::ListValue { - elements: elements.into_iter().map(Into::into).collect(), + elements: elements + .into_iter() + .map(TryInto::try_into) + .collect::>()?, }), SchemaValue::FixedList { elements } => { ValueBody::FixedListValue(proto::FixedListValue { - elements: elements.into_iter().map(Into::into).collect(), + elements: elements + .into_iter() + .map(TryInto::try_into) + .collect::>()?, }) } SchemaValue::Map { entries } => ValueBody::MapValue(proto::MapValue { entries: entries .into_iter() - .map(|(k, v)| proto::MapEntry { - key: Some(k.into()), - value: Some(v.into()), + .map(|(k, v)| { + Ok(proto::MapEntry { + key: Some(k.try_into()?), + value: Some(v.try_into()?), + }) }) - .collect(), + .collect::>()?, }), SchemaValue::Option { inner } => ValueBody::OptionValue(Box::new(proto::OptionValue { - inner: inner.map(|value| value_to_boxed_proto(*value)), + inner: inner + .map(|value| value_to_boxed_proto(*value)) + .transpose()?, })), SchemaValue::Result(r) => ValueBody::ResultValue(Box::new(proto::ResultValue { result: Some(match r { ResultValuePayload::Ok { value } => match value { - Some(v) => ResultBody::Ok(value_to_boxed_proto(*v)), + Some(v) => ResultBody::Ok(value_to_boxed_proto(*v)?), None => ResultBody::OkUnit(ProtoEmpty {}), }, ResultValuePayload::Err { value } => match value { - Some(v) => ResultBody::Err(value_to_boxed_proto(*v)), + Some(v) => ResultBody::Err(value_to_boxed_proto(*v)?), None => ResultBody::ErrUnit(ProtoEmpty {}), }, }), @@ -985,8 +1009,9 @@ impl From for proto::SchemaValue { SchemaValue::Quantity(q) => ValueBody::QuantityValue(q.into()), SchemaValue::Union(u) => ValueBody::UnionValue(Box::new(proto::UnionValue { tag: u.tag, - body: Some(value_to_boxed_proto(*u.body)), + body: Some(value_to_boxed_proto(*u.body)?), })), + #[cfg(not(all(feature = "guest", not(feature = "host"))))] SchemaValue::Secret(s) => ValueBody::SecretValue(proto::SecretValue { secret_id: Some(s.secret_id.into()), config_key: s.config_key.map(|items| proto::StringList { items }), @@ -994,6 +1019,7 @@ impl From for proto::SchemaValue { resolved_at: Some(datetime_to_proto(s.resolved_at)), category: s.category, }), + #[cfg(not(all(feature = "guest", not(feature = "host"))))] SchemaValue::QuotaToken(q) => ValueBody::QuotaTokenValue(proto::QuotaTokenValue { environment_id: Some(q.environment_id.uuid.into()), resource_name: q.resource_name, @@ -1001,8 +1027,25 @@ impl From for proto::SchemaValue { last_credit: q.last_credit, last_credit_at: Some(datetime_to_proto(q.last_credit_at)), }), + #[cfg(all(feature = "guest", not(feature = "host")))] + SchemaValue::Secret(_) => { + return Err( + "live secret handles cannot be converted to protobuf values".to_string() + ); + } + #[cfg(all(feature = "guest", not(feature = "host")))] + SchemaValue::QuotaToken(_) => { + return Err( + "live quota-token handles cannot be converted to protobuf values".to_string(), + ); + } + SchemaValue::Stream(_) => { + return Err( + "live schema value streams cannot be converted to protobuf values".to_string(), + ); + } }; - Self { value: Some(body) } + Ok(Self { value: Some(body) }) } } @@ -1109,6 +1152,7 @@ impl TryFrom for SchemaValue { body: req_box_value_from_proto(uv.body, "UnionValue.body")?, }) } + #[cfg(not(all(feature = "guest", not(feature = "host"))))] ValueBody::SecretValue(s) => SchemaValue::Secret(SecretValuePayload { secret_id: s .secret_id @@ -1122,6 +1166,7 @@ impl TryFrom for SchemaValue { )?, category: s.category, }), + #[cfg(not(all(feature = "guest", not(feature = "host"))))] ValueBody::QuotaTokenValue(q) => SchemaValue::QuotaToken(QuotaTokenValuePayload { environment_id: EnvironmentId::new( q.environment_id @@ -1137,18 +1182,34 @@ impl TryFrom for SchemaValue { })?, )?, }), + #[cfg(all(feature = "guest", not(feature = "host")))] + ValueBody::SecretValue(_) => { + return Err("protobuf values cannot create live secret handles".to_string()); + } + #[cfg(all(feature = "guest", not(feature = "host")))] + ValueBody::QuotaTokenValue(_) => { + return Err("protobuf values cannot create live quota-token handles".to_string()); + } + ValueBody::StreamReference(reference) => { + return Err(format!( + "live schema value stream reference {} cannot be decoded outside its session", + reference.stream_id + )); + } }; Ok(result) } } -impl From for proto::TypedSchemaValue { - fn from(value: TypedSchemaValue) -> Self { +impl TryFrom for proto::TypedSchemaValue { + type Error = String; + + fn try_from(value: TypedSchemaValue) -> Result { let (graph, val) = value.into_parts(); - Self { + Ok(Self { graph: Some(graph.into()), - value: Some(val.into()), - } + value: Some(val.try_into()?), + }) } } @@ -1173,6 +1234,35 @@ mod tests { use super::*; use test_r::test; + #[cfg(all(feature = "host", not(feature = "guest")))] + #[test] + fn generic_protobuf_encoding_rejects_live_streams() { + let value = SchemaValue::Stream(crate::schema::SchemaValueStream::from_host_endpoint(())); + + let error = proto::SchemaValue::try_from(value).unwrap_err(); + + assert_eq!( + error, + "live schema value streams cannot be converted to protobuf values" + ); + } + + #[test] + fn generic_protobuf_decoding_rejects_live_stream_references() { + let value = proto::SchemaValue { + value: Some(ValueBody::StreamReference( + proto::SchemaValueStreamReference { stream_id: 42 }, + )), + }; + + let error = SchemaValue::try_from(value).unwrap_err(); + + assert_eq!( + error, + "live schema value stream reference 42 cannot be decoded outside its session" + ); + } + #[test] fn proto_secret_spec_defaults_inner_to_string_when_absent() { let proto = proto::SchemaType { diff --git a/golem-schema/src/schema/schema_value.rs b/golem-schema/src/schema/schema_value.rs index 4e6c4a0564..158b7736d6 100644 --- a/golem-schema/src/schema/schema_value.rs +++ b/golem-schema/src/schema/schema_value.rs @@ -14,6 +14,7 @@ use crate::model::EnvironmentId; use crate::schema::schema_type::QuantityValue; +use crate::schema::stream::SchemaValueStream; use chrono::{DateTime, Utc}; use golem_schema_derive::{FromSchema, IntoSchema}; use serde::{Deserialize, Serialize}; @@ -130,8 +131,21 @@ pub enum SchemaValue { Union(UnionValuePayload), // Capability nodes + #[cfg_attr(all(feature = "guest", not(feature = "host")), serde(skip))] + #[cfg_attr( + all(feature = "full", feature = "guest", not(feature = "host")), + transient + )] Secret(SecretVariantValue), + #[cfg_attr(all(feature = "guest", not(feature = "host")), serde(skip))] + #[cfg_attr( + all(feature = "full", feature = "guest", not(feature = "host")), + transient + )] QuotaToken(QuotaTokenVariantValue), + #[serde(skip)] + #[cfg_attr(feature = "full", transient)] + Stream(SchemaValueStream), } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, IntoSchema, FromSchema)] diff --git a/golem-schema/src/schema/stream.rs b/golem-schema/src/schema/stream.rs new file mode 100644 index 0000000000..8e2c9b5fca --- /dev/null +++ b/golem-schema/src/schema/stream.rs @@ -0,0 +1,347 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Affine native streams carried recursively by [`super::SchemaValue`]. + +use crate::schema::conversion::{FromSchema, FromSchemaError, IntoSchema, SchemaBuilder}; +use crate::schema::metadata::{MetadataEnvelope, TypeId}; +use crate::schema::schema_type::SchemaType; +use crate::schema::schema_value::SchemaValue; + +#[cfg(all(feature = "guest", not(feature = "host")))] +type RawSchemaValueStream = wit_bindgen::StreamReader; + +#[cfg(all(feature = "host", not(feature = "guest")))] +mod active { + use std::any::Any; + use std::sync::{Arc, Mutex}; + + type OpaqueEndpoint = Box; + + /// Host-side stream leaf. Clones share one take-once endpoint; they never + /// duplicate a live stream. + /// + /// The endpoint is deliberately type-erased here. The schema model owns + /// affine transfer semantics, while the embedding runtime owns the relay + /// protocol used to connect Store-local Component Model streams. + #[derive(Clone)] + pub struct SchemaValueStream { + inner: Arc>>, + } + + impl SchemaValueStream { + #[doc(hidden)] + pub fn from_host_endpoint(endpoint: impl Any + Send) -> Self { + Self { + inner: Arc::new(Mutex::new(Some(Box::new(endpoint)))), + } + } + + #[doc(hidden)] + pub fn take_host_endpoint(&self) -> Result { + let endpoint = self + .inner + .lock() + .expect("schema value stream mutex poisoned") + .take() + .ok_or_else(|| "schema value stream was already transferred".to_string())?; + endpoint + .downcast::() + .map(|endpoint| *endpoint) + .map_err(|_| { + "schema value stream endpoint belongs to an incompatible runtime".to_string() + }) + } + + #[doc(hidden)] + pub fn with_host_endpoint( + &self, + f: impl FnOnce(&T) -> R, + ) -> Result { + let endpoint = self + .inner + .lock() + .expect("schema value stream mutex poisoned"); + let endpoint = endpoint + .as_ref() + .ok_or_else(|| "schema value stream was already transferred".to_string())?; + endpoint.downcast_ref::().map(f).ok_or_else(|| { + "schema value stream endpoint belongs to an incompatible runtime".to_string() + }) + } + + #[doc(hidden)] + pub fn take_for_transfer(&self) -> Option { + self.inner + .lock() + .expect("schema value stream mutex poisoned") + .take() + .map(|endpoint| Self { + inner: Arc::new(Mutex::new(Some(endpoint))), + }) + } + + #[doc(hidden)] + pub fn is_present(&self) -> bool { + self.inner + .lock() + .expect("schema value stream mutex poisoned") + .is_some() + } + + #[doc(hidden)] + pub fn cell_id(&self) -> *const () { + Arc::as_ptr(&self.inner).cast() + } + } + + impl std::fmt::Debug for SchemaValueStream { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("SchemaValueStream") + .field(&if self.is_present() { + "present" + } else { + "consumed" + }) + .finish() + } + } + + impl PartialEq for SchemaValueStream { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.inner, &other.inner) + } + } + + /// Resource-table representation of `golem:core/types.schema-value-stream`. + /// It owns one affine runtime endpoint. Component Model readers are + /// created or consumed only by the embedding runtime while it has access + /// to the endpoint's Store. + pub struct SchemaValueStreamHandleRep { + stream: SchemaValueStream, + } + + impl SchemaValueStreamHandleRep { + #[doc(hidden)] + pub fn new(stream: SchemaValueStream) -> Self { + Self { stream } + } + + #[doc(hidden)] + pub fn into_stream(self) -> SchemaValueStream { + self.stream + } + } +} + +#[cfg(all(feature = "guest", not(feature = "host")))] +mod active { + use super::RawSchemaValueStream; + use crate::schema::wit::wire; + use std::sync::{Arc, Mutex, MutexGuard}; + + enum State { + Wrapped(wire::SchemaValueStream), + Native(RawSchemaValueStream), + } + + /// Guest-side stream leaf. A leaf can hold the recursive WIT resource or + /// the native reader obtained from it. Clones share one take-once state. + #[derive(Clone)] + pub struct SchemaValueStream { + inner: Arc>>, + } + + impl SchemaValueStream { + fn state(&self) -> MutexGuard<'_, Option> { + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + } + + #[doc(hidden)] + pub fn from_native(reader: RawSchemaValueStream) -> Self { + Self { + inner: Arc::new(Mutex::new(Some(State::Native(reader)))), + } + } + + #[doc(hidden)] + pub fn from_wrapped(stream: wire::SchemaValueStream) -> Self { + Self { + inner: Arc::new(Mutex::new(Some(State::Wrapped(stream)))), + } + } + + #[doc(hidden)] + pub fn is_present(&self) -> bool { + self.state().is_some() + } + + #[doc(hidden)] + pub fn is_wrapped(&self) -> bool { + matches!(&*self.state(), Some(State::Wrapped(_))) + } + + #[doc(hidden)] + pub fn cell_id(&self) -> *const () { + Arc::as_ptr(&self.inner).cast() + } + + #[doc(hidden)] + pub fn take_wrapped(&self) -> Option { + let mut state = self.state(); + match state.take() { + Some(State::Wrapped(stream)) => Some(stream), + Some(native @ State::Native(_)) => { + *state = Some(native); + None + } + None => None, + } + } + + #[doc(hidden)] + pub async fn take_wrapped_async(&self) -> Result { + let state = self + .state() + .take() + .ok_or_else(|| "schema value stream was already transferred".to_string())?; + Ok(match state { + State::Wrapped(stream) => stream, + State::Native(reader) => wire::SchemaValueStream::wrap(reader).await, + }) + } + + #[doc(hidden)] + pub async fn ensure_wrapped(&self) -> Result<(), String> { + let state = self + .state() + .take() + .ok_or_else(|| "schema value stream was already transferred".to_string())?; + *self.state() = Some(match state { + State::Wrapped(stream) => State::Wrapped(stream), + State::Native(reader) => { + State::Wrapped(wire::SchemaValueStream::wrap(reader).await) + } + }); + Ok(()) + } + + #[doc(hidden)] + pub async fn take_native(self) -> Result { + let state = self + .state() + .take() + .ok_or_else(|| "schema value stream was already transferred".to_string())?; + Ok(match state { + State::Wrapped(stream) => wire::SchemaValueStream::unwrap(stream).await, + State::Native(reader) => reader, + }) + } + + #[doc(hidden)] + pub async fn next_wire(&self) -> Result, String> { + let state = self.state().take().ok_or_else(|| { + "schema value stream is already in use or was transferred".to_string() + })?; + let mut reader = match state { + State::Wrapped(stream) => wire::SchemaValueStream::unwrap(stream).await, + State::Native(reader) => reader, + }; + let item = reader.next().await; + *self.state() = Some(State::Native(reader)); + Ok(item) + } + } + + impl std::fmt::Debug for SchemaValueStream { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let state = match &*self.state() { + Some(State::Wrapped(_)) => "wrapped", + Some(State::Native(_)) => "native", + None => "consumed", + }; + f.debug_tuple("SchemaValueStream").field(&state).finish() + } + } + + impl PartialEq for SchemaValueStream { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.inner, &other.inner) + } + } +} + +#[cfg(any( + all(feature = "host", not(feature = "guest")), + all(feature = "guest", not(feature = "host")) +))] +pub use active::SchemaValueStream; + +#[cfg(all(feature = "host", not(feature = "guest")))] +pub use active::SchemaValueStreamHandleRep; + +#[cfg(not(any( + all(feature = "host", not(feature = "guest")), + all(feature = "guest", not(feature = "host")) +)))] +#[derive(Clone, Debug, PartialEq)] +pub struct SchemaValueStream(()); + +impl serde::Serialize for SchemaValueStream { + fn serialize(&self, _serializer: S) -> Result { + Err(serde::ser::Error::custom( + "live schema value streams cannot be serialized", + )) + } +} + +impl<'de> serde::Deserialize<'de> for SchemaValueStream { + fn deserialize>(_deserializer: D) -> Result { + Err(serde::de::Error::custom( + "live schema value streams cannot be deserialized", + )) + } +} + +impl IntoSchema for SchemaValueStream { + fn type_id() -> TypeId { + TypeId::new("golem.core.SchemaValueStream") + } + + fn register_in(_builder: &mut SchemaBuilder) -> SchemaType { + SchemaType::Stream { + inner: None, + metadata: MetadataEnvelope::default(), + } + } + + fn to_value(&self) -> SchemaValue { + SchemaValue::Stream(self.clone()) + } +} + +impl FromSchema for SchemaValueStream { + fn from_value(value: &SchemaValue) -> Result { + match value { + SchemaValue::Stream(stream) => Ok(stream.clone()), + other => Err(FromSchemaError::shape_mismatch( + "stream", + format!("{other:?}"), + "SchemaValueStream", + )), + } + } +} diff --git a/golem-schema/src/schema/tool/canonical.rs b/golem-schema/src/schema/tool/canonical.rs index 225234b79d..9870f95f49 100644 --- a/golem-schema/src/schema/tool/canonical.rs +++ b/golem-schema/src/schema/tool/canonical.rs @@ -53,6 +53,7 @@ use std::collections::BTreeSet; pub struct CanonicalInputField { pub name: String, pub aliases: Vec, + pub short: Option, pub type_: SchemaType, } @@ -70,6 +71,7 @@ pub struct CanonicalInputModel { pub struct CanonicalInputValue { pub name: String, pub aliases: Vec, + pub short: Option, pub type_: SchemaType, pub value: SchemaValue, } @@ -156,6 +158,7 @@ impl CanonicalInputModel { .map(|(field, value)| CanonicalInputValue { name: field.name, aliases: field.aliases, + short: field.short, type_: field.type_, value, }) @@ -324,6 +327,7 @@ impl Tool { Some(CanonicalInputField { name: option.long.clone(), aliases: option.aliases.clone(), + short: option.short, type_: option_collected_type(&option.shape), }) } @@ -332,6 +336,7 @@ impl Tool { Some(CanonicalInputField { name: flag.long.clone(), aliases: flag.aliases.clone(), + short: flag.short, type_: flag_type(flag), }) } @@ -340,6 +345,7 @@ impl Tool { Some(CanonicalInputField { name: positional.name.clone(), aliases: Vec::new(), + short: None, type_: positional.type_.clone(), }) } @@ -348,6 +354,7 @@ impl Tool { Some(CanonicalInputField { name: tail.name.clone(), aliases: Vec::new(), + short: None, type_: tail_collected_type(tail), }) } @@ -356,6 +363,7 @@ impl Tool { Some(CanonicalInputField { name: option.long.clone(), aliases: option.aliases.clone(), + short: option.short, type_: option_collected_type(&option.shape), }) } @@ -364,6 +372,7 @@ impl Tool { Some(CanonicalInputField { name: flag.long.clone(), aliases: flag.aliases.clone(), + short: flag.short, type_: flag_type(flag), }) } @@ -868,6 +877,23 @@ mod tests { ); } + #[test] + fn canonical_input_fields_preserve_short_options() { + let mut tool = grep_tool(); + tool.commands.nodes[0].globals.options[0].short = Some('c'); + tool.commands.nodes[0].globals.flags[0].short = Some('i'); + tool.commands.nodes[0].body.as_mut().unwrap().options[1].short = Some('n'); + tool.commands.nodes[0].body.as_mut().unwrap().flags[0].short = Some('v'); + + let fields = tool.canonical_input_fields(0); + let shorts: Vec> = fields.iter().map(|field| field.short).collect(); + + assert_eq!( + shorts, + vec![Some('c'), Some('i'), None, None, None, Some('n'), Some('v'),] + ); + } + #[test] fn globals_are_effective_on_subcommand() { let tool = grep_tool(); diff --git a/golem-schema/src/schema/tool/protobuf.rs b/golem-schema/src/schema/tool/protobuf.rs index 170c265341..4dfaac2cfe 100644 --- a/golem-schema/src/schema/tool/protobuf.rs +++ b/golem-schema/src/schema/tool/protobuf.rs @@ -28,6 +28,12 @@ fn decode_char(value: u32, field: &str) -> Result { char::from_u32(value).ok_or_else(|| format!("Invalid Unicode scalar in {field}: {value}")) } +fn encode_static_value(value: SchemaValue) -> golem_api_grpc::proto::golem::schema::SchemaValue { + value + .try_into() + .expect("static tool values cannot contain live streams") +} + impl From for proto::Tool { fn from(value: Tool) -> Self { Self { @@ -231,7 +237,7 @@ impl From for proto::Positional { doc: Some(value.doc.into()), value_name: value.value_name, r#type: Some(value.type_.into()), - default: value.default.map(Into::into), + default: value.default.map(encode_static_value), required: value.required, accepts_stdio: value.accepts_stdio, } @@ -297,7 +303,7 @@ impl From for proto::OptionSpec { doc: Some(value.doc.into()), value_name: value.value_name, shape: Some(value.shape.into()), - default: value.default.map(Into::into), + default: value.default.map(encode_static_value), required: value.required, env_var: value.env_var, } @@ -536,7 +542,7 @@ impl From for proto::ValueIsRef { fn from(value: ValueIsRef) -> Self { Self { name: value.name, - value: Some(value.value.into()), + value: Some(encode_static_value(value.value)), } } } diff --git a/golem-schema/src/schema/validation/value.rs b/golem-schema/src/schema/validation/value.rs index 8e24cf6398..eb367ec869 100644 --- a/golem-schema/src/schema/validation/value.rs +++ b/golem-schema/src/schema/validation/value.rs @@ -557,6 +557,7 @@ fn shape_name(value: &SchemaValue) -> &'static str { SchemaValue::Union(_) => "union", SchemaValue::Secret(_) => "secret", SchemaValue::QuotaToken(_) => "quota-token", + SchemaValue::Stream(_) => "stream", } } @@ -813,6 +814,7 @@ fn check<'a>( #[cfg(all(feature = "guest", not(feature = "host")))] let _ = (spec, payload); } + (SchemaType::Stream { .. }, SchemaValue::Stream(_)) => {} (SchemaType::Record { fields, .. }, SchemaValue::Record { fields: vs }) => { if fields.len() != vs.len() { diff --git a/golem-schema/src/schema/wit/decode.rs b/golem-schema/src/schema/wit/decode.rs index 8343b70300..674299b2a5 100644 --- a/golem-schema/src/schema/wit/decode.rs +++ b/golem-schema/src/schema/wit/decode.rs @@ -43,6 +43,11 @@ type WireSecretHandle = wasmtime::component::Resource; #[cfg(all(feature = "guest", not(feature = "host")))] type WireSecretHandle = wire::Secret; +#[cfg(all(feature = "host", not(feature = "guest")))] +type WireStreamHandle = wasmtime::component::Resource; +#[cfg(all(feature = "guest", not(feature = "host")))] +type WireStreamHandle = wire::SchemaValueStream; + /// Errors that can occur while decoding the flat wire form into the /// recursive in-memory representation. #[derive(Debug, Clone, PartialEq, Eq)] @@ -78,6 +83,8 @@ pub enum DecodeError { /// An owned `secret` handle was present in the value tree but never reached /// from the root. UnconsumedSecretHandle(wire::ValueNodeIndex), + /// An owned stream wrapper was present but never reached from the root. + UnconsumedStream(wire::ValueNodeIndex), /// The host `QuotaTokenResolver` failed to snapshot an owned handle. QuotaResolver(String), /// The host `SecretResolver` failed to snapshot an owned handle. @@ -92,6 +99,12 @@ pub enum DecodeError { /// does not permit secret transport. The handle is dropped from the resource /// table before this error is returned, so nothing leaks. SecretNotPermitted(wire::ValueNodeIndex), + /// A stream wrapper was encountered at a non-stream-aware host boundary. + StreamRequiresStore(wire::ValueNodeIndex), + /// A live stream was encountered at a materializing boundary. + StreamNotPermitted(wire::ValueNodeIndex), + /// The host failed to transfer a schema-value-stream resource into a native reader. + StreamResolver(String), } impl Display for DecodeError { @@ -130,6 +143,9 @@ impl Display for DecodeError { DecodeError::UnconsumedSecretHandle(i) => { write!(f, "secret handle not referenced from the root: {i}") } + DecodeError::UnconsumedStream(i) => { + write!(f, "schema value stream not referenced from the root: {i}") + } DecodeError::QuotaResolver(msg) => { write!(f, "quota-token handle could not be resolved: {msg}") } @@ -142,6 +158,21 @@ impl Display for DecodeError { DecodeError::SecretNotPermitted(i) => { write!(f, "secret handle not permitted at this boundary: {i}") } + DecodeError::StreamRequiresStore(i) => { + write!( + f, + "schema value stream requires a stream-aware store boundary: {i}" + ) + } + DecodeError::StreamNotPermitted(i) => { + write!( + f, + "live schema value stream not permitted at this boundary: {i}" + ) + } + DecodeError::StreamResolver(message) => { + write!(f, "schema value stream could not be transferred: {message}") + } } } } @@ -240,6 +271,7 @@ pub fn decode_value(wire_tree: wire::SchemaValueTree) -> Result Err(e), }; @@ -260,6 +292,9 @@ pub fn decode_value(wire_tree: wire::SchemaValueTree) -> Result { + leaked.get_or_insert(DecodeError::UnconsumedStream(i as wire::ValueNodeIndex)); + } other => *slot = other, } } @@ -709,27 +744,30 @@ fn decode_value_node( wire::SchemaValueNode::QuotaTokenHandle(_) => { return Err(DecodeError::QuotaTokenRequiresResolver); } + wire::SchemaValueNode::StreamValue(_) => { + return Err(DecodeError::StreamRequiresStore(0)); + } }; Ok(out) } -/// Decode an owned value tree, lifting each `quota-token-handle` node into a -/// trusted [`SchemaValue::QuotaToken`] snapshot via the supplied resolver. +/// Decode an owned value tree, lifting each resource node into its native +/// [`SchemaValue`] representation via the supplied resolver. /// /// The tree is consumed because the owned handles it carries are affine: each /// must be moved out and snapshotted exactly once. Any node referenced more /// than once (or forming a cycle) is rejected with /// [`DecodeError::AliasedValueNode`]. /// -/// Every `own` handle present in the tree has already been -/// transferred to the host, so each must be consumed exactly once. After -/// decoding, any handle that was never reached from the root is released through -/// [`super::QuotaTokenResolver::drop_handle`] (so none leak from the table) and, -/// on an otherwise successful decode, the tree is rejected as malformed with -/// [`DecodeError::UnconsumedQuotaTokenHandle`]. The same cleanup runs when -/// decoding fails partway through. +/// Every owned resource present in the tree has already been transferred to +/// the host, so each must be consumed exactly once. After decoding, any handle +/// that was never reached from the root is released through its resolver (so +/// none leak from the table) and the tree is rejected as malformed. The same +/// cleanup runs when decoding fails partway through. #[cfg(all(feature = "host", not(feature = "guest")))] -pub fn decode_value_with( +pub fn decode_value_with< + R: super::QuotaTokenResolver + super::SecretResolver + super::SchemaValueStreamResolver, +>( wire_tree: wire::SchemaValueTree, resolver: &mut R, ) -> Result { @@ -737,55 +775,67 @@ pub fn decode_value_with( let root = wire_tree.root; let mut slots: Vec> = wire_tree.value_nodes.into_iter().map(Some).collect(); - // Validate the whole tree before snapshotting any owned handle, so a - // malformed sibling cannot cause an already-snapshotted token to be - // discarded. After a successful preflight the only fallible step left is the - // snapshot itself. + // Validate the whole tree before resolving any owned handle, so a + // malformed sibling cannot cause an already-resolved resource to be + // discarded. After a successful preflight the only fallible steps left are + // the resolver calls themselves. let result = match preflight_owned_value_tree(&slots, root) { Ok(()) => decode_owned_at( &mut slots, root, &mut |handle| { - let mut resolver = resolver.borrow_mut(); resolver + .borrow_mut() .snapshot_handle(handle) - .map_err(|e| DecodeError::QuotaResolver(e.to_string())) + .map_err(|error| DecodeError::QuotaResolver(error.to_string())) }, &mut |handle| { - let mut resolver = resolver.borrow_mut(); resolver + .borrow_mut() .snapshot_secret_handle(handle) - .map_err(|e| DecodeError::SecretResolver(e.to_string())) + .map_err(|error| DecodeError::SecretResolver(error.to_string())) + }, + &mut |handle| { + resolver + .borrow_mut() + .stream_from_handle(handle) + .map_err(|error| DecodeError::StreamResolver(error.to_string())) }, ), - Err(e) => Err(e), + Err(error) => Err(error), }; - // Drop every handle that was not consumed while walking the tree, regardless - // of success or failure, so no owned resource leaks from the table. - let mut leaked: Option = None; + // Drop every handle that was not consumed while walking the tree, + // regardless of success or failure, so no owned resource leaks from the + // table. + let mut unconsumed = None; let mut resolver = resolver.borrow_mut(); - for (i, slot) in slots.iter_mut().enumerate() { + for (index, slot) in slots.iter_mut().enumerate() { match slot.take() { Some(wire::SchemaValueNode::QuotaTokenHandle(handle)) => { resolver.drop_handle(handle); - leaked.get_or_insert(DecodeError::UnconsumedQuotaTokenHandle( - i as wire::ValueNodeIndex, + unconsumed.get_or_insert(DecodeError::UnconsumedQuotaTokenHandle( + index as wire::ValueNodeIndex, )); } Some(wire::SchemaValueNode::SecretValue(handle)) => { super::SecretResolver::drop_secret_handle(&mut **resolver, handle); - leaked.get_or_insert(DecodeError::UnconsumedSecretHandle( - i as wire::ValueNodeIndex, + unconsumed.get_or_insert(DecodeError::UnconsumedSecretHandle( + index as wire::ValueNodeIndex, )); } + Some(wire::SchemaValueNode::StreamValue(handle)) => { + resolver.drop_stream_handle(handle); + unconsumed + .get_or_insert(DecodeError::UnconsumedStream(index as wire::ValueNodeIndex)); + } other => *slot = other, } } match result { - Ok(value) => leaked.map_or(Ok(value), Err), - Err(err) => Err(err), + Ok(value) => unconsumed.map_or(Ok(value), Err), + Err(error) => Err(error), } } @@ -942,6 +992,9 @@ fn preflight_owned_value_tree( i as wire::ValueNodeIndex, )); } + Some(wire::SchemaValueNode::StreamValue(_)) if !reached[i] => { + return Err(DecodeError::UnconsumedStream(i as wire::ValueNodeIndex)); + } _ => {} } } @@ -1007,6 +1060,7 @@ fn preflight_owned_at( datetime_from_wire(d)?; } wire::SchemaValueNode::SecretValue(_) => {} + wire::SchemaValueNode::StreamValue(_) => {} // All remaining node kinds are leaves with no child indices and no // extra decode-time validation. Quota handles are leaves too; their // reachability is checked by the caller after the walk. @@ -1024,6 +1078,9 @@ fn reject_handles_in_pure_value_tree(wire_tree: &wire::SchemaValueTree) -> Resul wire::SchemaValueNode::QuotaTokenHandle(_) => { return Err(DecodeError::QuotaTokenRequiresResolver); } + wire::SchemaValueNode::StreamValue(_) => { + return Err(DecodeError::StreamRequiresStore(0)); + } _ => {} } } @@ -1044,6 +1101,9 @@ fn decode_owned_at( idx: wire::ValueNodeIndex, lift_quota: &mut dyn FnMut(WireQuotaHandle) -> Result, lift_secret: &mut dyn FnMut(WireSecretHandle) -> Result, + lift_stream: &mut dyn FnMut( + WireStreamHandle, + ) -> Result, ) -> Result { let pos = usize_index_v(idx)?; let node = slots @@ -1068,7 +1128,13 @@ fn decode_owned_at( wire::SchemaValueNode::RecordValue(fields) => { let mut decoded = Vec::with_capacity(fields.len()); for i in fields { - decoded.push(decode_owned_at(slots, i, lift_quota, lift_secret)?); + decoded.push(decode_owned_at( + slots, + i, + lift_quota, + lift_secret, + lift_stream, + )?); } SchemaValue::Record { fields: decoded } } @@ -1079,6 +1145,7 @@ fn decode_owned_at( i, lift_quota, lift_secret, + lift_stream, )?)), None => None, }; @@ -1092,29 +1159,47 @@ fn decode_owned_at( wire::SchemaValueNode::TupleValue(elements) => { let mut decoded = Vec::with_capacity(elements.len()); for i in elements { - decoded.push(decode_owned_at(slots, i, lift_quota, lift_secret)?); + decoded.push(decode_owned_at( + slots, + i, + lift_quota, + lift_secret, + lift_stream, + )?); } SchemaValue::Tuple { elements: decoded } } wire::SchemaValueNode::ListValue(elements) => { let mut decoded = Vec::with_capacity(elements.len()); for i in elements { - decoded.push(decode_owned_at(slots, i, lift_quota, lift_secret)?); + decoded.push(decode_owned_at( + slots, + i, + lift_quota, + lift_secret, + lift_stream, + )?); } SchemaValue::List { elements: decoded } } wire::SchemaValueNode::FixedListValue(elements) => { let mut decoded = Vec::with_capacity(elements.len()); for i in elements { - decoded.push(decode_owned_at(slots, i, lift_quota, lift_secret)?); + decoded.push(decode_owned_at( + slots, + i, + lift_quota, + lift_secret, + lift_stream, + )?); } SchemaValue::FixedList { elements: decoded } } wire::SchemaValueNode::MapValue(entries) => { let mut decoded = Vec::with_capacity(entries.len()); for e in entries { - let key = decode_owned_at(slots, e.key, lift_quota, lift_secret)?; - let value = decode_owned_at(slots, e.value, lift_quota, lift_secret)?; + let key = decode_owned_at(slots, e.key, lift_quota, lift_secret, lift_stream)?; + let value = decode_owned_at(slots, e.value, lift_quota, lift_secret, lift_stream)?; decoded.push((key, value)); } SchemaValue::Map { entries: decoded } @@ -1126,6 +1211,7 @@ fn decode_owned_at( i, lift_quota, lift_secret, + lift_stream, )?)), None => None, }, @@ -1139,6 +1225,7 @@ fn decode_owned_at( i, lift_quota, lift_secret, + lift_stream, )?)), None => None, }, @@ -1150,6 +1237,7 @@ fn decode_owned_at( i, lift_quota, lift_secret, + lift_stream, )?)), None => None, }, @@ -1180,12 +1268,19 @@ fn decode_owned_at( }), wire::SchemaValueNode::UnionValue(p) => SchemaValue::Union(UnionValuePayload { tag: p.tag, - body: Box::new(decode_owned_at(slots, p.body, lift_quota, lift_secret)?), + body: Box::new(decode_owned_at( + slots, + p.body, + lift_quota, + lift_secret, + lift_stream, + )?), }), wire::SchemaValueNode::SecretValue(handle) => SchemaValue::Secret(lift_secret(handle)?), wire::SchemaValueNode::QuotaTokenHandle(handle) => { SchemaValue::QuotaToken(lift_quota(handle)?) } + wire::SchemaValueNode::StreamValue(stream) => SchemaValue::Stream(lift_stream(stream)?), }; Ok(out) } diff --git a/golem-schema/src/schema/wit/encode.rs b/golem-schema/src/schema/wit/encode.rs index ccc36b5059..471894273a 100644 --- a/golem-schema/src/schema/wit/encode.rs +++ b/golem-schema/src/schema/wit/encode.rs @@ -64,6 +64,17 @@ pub enum EncodeError { /// (Guest) The same owned secret handle appeared more than once in a single /// value tree. AliasedSecretHandle, + /// A live stream was used at a materializing boundary rather than a + /// stream-aware invocation boundary. + StreamNotTransportable, + /// A stream endpoint was already transferred by an earlier encode. + StreamAlreadyConsumed, + /// The same affine stream leaf appeared more than once in one value tree. + AliasedStream, + /// A guest-native reader must first be wrapped asynchronously. + StreamRequiresAsyncEncoding, + /// The host failed to transfer a native reader into a schema-value-stream resource. + StreamResolver(String), } impl Display for EncodeError { @@ -103,6 +114,26 @@ impl Display for EncodeError { "the same secret handle appeared more than once in one value tree" ) } + EncodeError::StreamNotTransportable => write!( + f, + "live streams cannot cross this materializing value boundary" + ), + EncodeError::StreamAlreadyConsumed => { + write!(f, "schema value stream was already transferred") + } + EncodeError::AliasedStream => { + write!( + f, + "the same affine stream appeared more than once in one value tree" + ) + } + EncodeError::StreamRequiresAsyncEncoding => write!( + f, + "a native schema value stream requires the asynchronous guest encoder" + ), + EncodeError::StreamResolver(message) => { + write!(f, "schema value stream could not be transferred: {message}") + } } } } @@ -178,6 +209,7 @@ pub fn encode_value(value: &SchemaValue) -> Result Result Result { + preflight_guest_handles(value)?; + let mut streams = Vec::new(); + collect_streams(value, &mut streams); + for stream in streams { + stream + .ensure_wrapped() + .await + .map_err(|_| EncodeError::StreamAlreadyConsumed)?; + } + encode_value(value) +} + /// Walk a value tree and verify that every quota-token handle is still present /// and unique, without taking any handle. Returns an error if a handle was /// already consumed or the same handle appears more than once. @@ -221,6 +276,7 @@ fn preflight_guest_handles(value: &SchemaValue) -> Result<(), EncodeError> { value: &SchemaValue, seen_quota: &mut std::collections::HashSet<*const ()>, seen_secret: &mut std::collections::HashSet<*const ()>, + seen_stream: &mut std::collections::HashSet<*const ()>, ) -> Result<(), EncodeError> { match value { SchemaValue::QuotaToken(handle) => { @@ -241,9 +297,18 @@ fn preflight_guest_handles(value: &SchemaValue) -> Result<(), EncodeError> { } Ok(()) } + SchemaValue::Stream(stream) => { + if !stream.is_present() { + return Err(EncodeError::StreamAlreadyConsumed); + } + if !seen_stream.insert(stream.cell_id()) { + return Err(EncodeError::AliasedStream); + } + Ok(()) + } SchemaValue::Record { fields } => { for f in fields { - walk(f, seen_quota, seen_secret)?; + walk(f, seen_quota, seen_secret, seen_stream)?; } Ok(()) } @@ -251,26 +316,26 @@ fn preflight_guest_handles(value: &SchemaValue) -> Result<(), EncodeError> { | SchemaValue::List { elements } | SchemaValue::FixedList { elements } => { for e in elements { - walk(e, seen_quota, seen_secret)?; + walk(e, seen_quota, seen_secret, seen_stream)?; } Ok(()) } SchemaValue::Variant(p) => { if let Some(payload) = &p.payload { - walk(payload, seen_quota, seen_secret)?; + walk(payload, seen_quota, seen_secret, seen_stream)?; } Ok(()) } SchemaValue::Map { entries } => { for (k, v) in entries { - walk(k, seen_quota, seen_secret)?; - walk(v, seen_quota, seen_secret)?; + walk(k, seen_quota, seen_secret, seen_stream)?; + walk(v, seen_quota, seen_secret, seen_stream)?; } Ok(()) } SchemaValue::Option { inner } => { if let Some(inner) = inner { - walk(inner, seen_quota, seen_secret)?; + walk(inner, seen_quota, seen_secret, seen_stream)?; } Ok(()) } @@ -279,18 +344,67 @@ fn preflight_guest_handles(value: &SchemaValue) -> Result<(), EncodeError> { ResultValuePayload::Ok { value } | ResultValuePayload::Err { value } => value, }; if let Some(inner) = inner { - walk(inner, seen_quota, seen_secret)?; + walk(inner, seen_quota, seen_secret, seen_stream)?; } Ok(()) } - SchemaValue::Union(p) => walk(&p.body, seen_quota, seen_secret), + SchemaValue::Union(p) => walk(&p.body, seen_quota, seen_secret, seen_stream), _ => Ok(()), } } let mut seen_quota = std::collections::HashSet::new(); let mut seen_secret = std::collections::HashSet::new(); - walk(value, &mut seen_quota, &mut seen_secret) + let mut seen_stream = std::collections::HashSet::new(); + walk(value, &mut seen_quota, &mut seen_secret, &mut seen_stream) +} + +#[cfg(all(feature = "guest", not(feature = "host")))] +fn collect_streams<'a>( + value: &'a SchemaValue, + streams: &mut Vec<&'a crate::schema::SchemaValueStream>, +) { + match value { + SchemaValue::Stream(stream) => streams.push(stream), + SchemaValue::Record { fields } => { + for value in fields { + collect_streams(value, streams); + } + } + SchemaValue::Tuple { elements } + | SchemaValue::List { elements } + | SchemaValue::FixedList { elements } => { + for value in elements { + collect_streams(value, streams); + } + } + SchemaValue::Variant(payload) => { + if let Some(value) = &payload.payload { + collect_streams(value, streams); + } + } + SchemaValue::Map { entries } => { + for (key, value) in entries { + collect_streams(key, streams); + collect_streams(value, streams); + } + } + SchemaValue::Option { inner } => { + if let Some(value) = inner { + collect_streams(value, streams); + } + } + SchemaValue::Result(payload) => { + let value = match payload { + ResultValuePayload::Ok { value } | ResultValuePayload::Err { value } => value, + }; + if let Some(value) = value { + collect_streams(value, streams); + } + } + SchemaValue::Union(payload) => collect_streams(&payload.body, streams), + _ => {} + } } /// Encode a value tree, turning each [`SchemaValue::QuotaToken`] snapshot into a @@ -328,6 +442,7 @@ pub fn encode_value_with( .map_err(|e| EncodeError::SecretResolver(e.to_string()))?; Ok(wire::SchemaValueNode::SecretValue(handle)) }, + &mut |_stream| Err(EncodeError::StreamNotTransportable), ); match root { Ok(root) => Ok(wire::SchemaValueTree { @@ -350,13 +465,134 @@ pub fn encode_value_with( } } +#[cfg(all(feature = "host", not(feature = "guest")))] +pub fn encode_value_with_streams< + R: super::QuotaTokenResolver + super::SecretResolver + super::SchemaValueStreamResolver, +>( + value: &SchemaValue, + resolver: &mut R, +) -> Result { + preflight_host_streams(value)?; + let resolver = std::cell::RefCell::new(resolver); + let mut ctx = ValueCtx::default(); + let root = ctx.encode( + value, + &mut |snapshot| { + let handle = resolver + .borrow_mut() + .handle_from_snapshot(snapshot) + .map_err(|error| EncodeError::QuotaResolver(error.to_string()))?; + Ok(wire::SchemaValueNode::QuotaTokenHandle(handle)) + }, + &mut |snapshot| { + let handle = resolver + .borrow_mut() + .secret_handle_from_snapshot(snapshot) + .map_err(|error| EncodeError::SecretResolver(error.to_string()))?; + Ok(wire::SchemaValueNode::SecretValue(handle)) + }, + &mut |stream| { + let stream = stream + .take_for_transfer() + .ok_or(EncodeError::StreamAlreadyConsumed)?; + let handle = resolver + .borrow_mut() + .handle_from_stream(stream) + .map_err(|error| EncodeError::StreamResolver(error.to_string()))?; + Ok(wire::SchemaValueNode::StreamValue(handle)) + }, + ); + match root { + Ok(root) => Ok(wire::SchemaValueTree { + value_nodes: ctx.value_nodes, + root, + }), + Err(error) => { + let mut resolver = resolver.borrow_mut(); + for node in ctx.value_nodes { + match node { + wire::SchemaValueNode::QuotaTokenHandle(handle) => resolver.drop_handle(handle), + wire::SchemaValueNode::SecretValue(handle) => { + super::SecretResolver::drop_secret_handle(&mut **resolver, handle) + } + wire::SchemaValueNode::StreamValue(handle) => { + resolver.drop_stream_handle(handle) + } + _ => {} + } + } + Err(error) + } + } +} + +#[cfg(all(feature = "host", not(feature = "guest")))] +fn preflight_host_streams(value: &SchemaValue) -> Result<(), EncodeError> { + fn walk( + value: &SchemaValue, + seen: &mut std::collections::HashSet<*const ()>, + ) -> Result<(), EncodeError> { + match value { + SchemaValue::Stream(stream) => { + if !stream.is_present() { + return Err(EncodeError::StreamAlreadyConsumed); + } + if !seen.insert(stream.cell_id()) { + return Err(EncodeError::AliasedStream); + } + } + SchemaValue::Record { fields } => { + for value in fields { + walk(value, seen)?; + } + } + SchemaValue::Tuple { elements } + | SchemaValue::List { elements } + | SchemaValue::FixedList { elements } => { + for value in elements { + walk(value, seen)?; + } + } + SchemaValue::Variant(payload) => { + if let Some(value) = &payload.payload { + walk(value, seen)?; + } + } + SchemaValue::Map { entries } => { + for (key, value) in entries { + walk(key, seen)?; + walk(value, seen)?; + } + } + SchemaValue::Option { inner: Some(value) } => walk(value, seen)?, + SchemaValue::Option { inner: None } => {} + SchemaValue::Result(payload) => { + let value = match payload { + ResultValuePayload::Ok { value } | ResultValuePayload::Err { value } => value, + }; + if let Some(value) = value { + walk(value, seen)?; + } + } + SchemaValue::Union(payload) => walk(&payload.body, seen)?, + _ => {} + } + Ok(()) + } + + walk(value, &mut std::collections::HashSet::new()) +} + fn encode_value_inner( value: &SchemaValue, quota: &mut dyn FnMut(&QuotaTokenVariantValue) -> Result, secret: &mut dyn FnMut(&SecretVariantValue) -> Result, + stream: &mut dyn FnMut( + &crate::schema::SchemaValueStream, + ) -> Result, ) -> Result { let mut ctx = ValueCtx::default(); - let root = ctx.encode(value, quota, secret)?; + let root = ctx.encode(value, quota, secret, stream)?; Ok(wire::SchemaValueTree { value_nodes: ctx.value_nodes, root, @@ -737,6 +973,9 @@ impl ValueCtx { &QuotaTokenVariantValue, ) -> Result, secret: &mut dyn FnMut(&SecretVariantValue) -> Result, + stream: &mut dyn FnMut( + &crate::schema::SchemaValueStream, + ) -> Result, ) -> Result { let node = match value { SchemaValue::Bool(b) => wire::SchemaValueNode::BoolValue(*b), @@ -755,13 +994,13 @@ impl ValueCtx { SchemaValue::Record { fields } => { let mut indices = Vec::with_capacity(fields.len()); for v in fields { - indices.push(self.encode(v, quota, secret)?); + indices.push(self.encode(v, quota, secret, stream)?); } wire::SchemaValueNode::RecordValue(indices) } SchemaValue::Variant(p) => { let payload = match &p.payload { - Some(v) => Some(self.encode(v, quota, secret)?), + Some(v) => Some(self.encode(v, quota, secret, stream)?), None => None, }; wire::SchemaValueNode::VariantValue(wire::VariantValuePayload { @@ -774,21 +1013,21 @@ impl ValueCtx { SchemaValue::Tuple { elements } => { let mut indices = Vec::with_capacity(elements.len()); for v in elements { - indices.push(self.encode(v, quota, secret)?); + indices.push(self.encode(v, quota, secret, stream)?); } wire::SchemaValueNode::TupleValue(indices) } SchemaValue::List { elements } => { let mut indices = Vec::with_capacity(elements.len()); for v in elements { - indices.push(self.encode(v, quota, secret)?); + indices.push(self.encode(v, quota, secret, stream)?); } wire::SchemaValueNode::ListValue(indices) } SchemaValue::FixedList { elements } => { let mut indices = Vec::with_capacity(elements.len()); for v in elements { - indices.push(self.encode(v, quota, secret)?); + indices.push(self.encode(v, quota, secret, stream)?); } wire::SchemaValueNode::FixedListValue(indices) } @@ -796,15 +1035,15 @@ impl ValueCtx { let mut encoded = Vec::with_capacity(entries.len()); for (k, v) in entries { encoded.push(wire::MapEntry { - key: self.encode(k, quota, secret)?, - value: self.encode(v, quota, secret)?, + key: self.encode(k, quota, secret, stream)?, + value: self.encode(v, quota, secret, stream)?, }); } wire::SchemaValueNode::MapValue(encoded) } SchemaValue::Option { inner } => { let inner = match inner { - Some(v) => Some(self.encode(v, quota, secret)?), + Some(v) => Some(self.encode(v, quota, secret, stream)?), None => None, }; wire::SchemaValueNode::OptionValue(inner) @@ -813,14 +1052,14 @@ impl ValueCtx { let payload = match p { ResultValuePayload::Ok { value } => { let v = match value { - Some(v) => Some(self.encode(v, quota, secret)?), + Some(v) => Some(self.encode(v, quota, secret, stream)?), None => None, }; wire::ResultValuePayload::OkValue(v) } ResultValuePayload::Err { value } => { let v = match value { - Some(v) => Some(self.encode(v, quota, secret)?), + Some(v) => Some(self.encode(v, quota, secret, stream)?), None => None, }; wire::ResultValuePayload::ErrValue(v) @@ -861,7 +1100,7 @@ impl ValueCtx { }) } SchemaValue::Union(p) => { - let body = self.encode(&p.body, quota, secret)?; + let body = self.encode(&p.body, quota, secret, stream)?; wire::SchemaValueNode::UnionValue(wire::UnionValuePayload { tag: p.tag.clone(), body, @@ -869,6 +1108,7 @@ impl ValueCtx { } SchemaValue::Secret(s) => secret(s)?, SchemaValue::QuotaToken(q) => quota(q)?, + SchemaValue::Stream(s) => stream(s)?, }; Ok(self.push(node)) } diff --git a/golem-schema/src/schema/wit/host.rs b/golem-schema/src/schema/wit/host.rs index e8eefe6d12..230fd1e971 100644 --- a/golem-schema/src/schema/wit/host.rs +++ b/golem-schema/src/schema/wit/host.rs @@ -19,7 +19,9 @@ pub mod generated { path: "wit", world: "golem-schema", imports: { - default: async | trappable, + "golem:core/types.[static]schema-value-stream.unwrap": async | store | trappable, + "golem:core/types.[static]schema-value-stream.wrap": async | store | trappable, + default: async | store, }, exports: { default: async }, require_store_data_send: true, @@ -28,6 +30,7 @@ pub mod generated { with: { "golem:core/types@2.0.0.quota-token": crate::schema::wit::QuotaTokenHandleRep, "golem:core/types@2.0.0.secret": crate::schema::wit::SecretHandleRep, + "golem:core/types@2.0.0.schema-value-stream": crate::schema::SchemaValueStreamHandleRep, }, }); } diff --git a/golem-schema/src/schema/wit/mod.rs b/golem-schema/src/schema/wit/mod.rs index 959ca00ca6..a019d28ce3 100644 --- a/golem-schema/src/schema/wit/mod.rs +++ b/golem-schema/src/schema/wit/mod.rs @@ -46,6 +46,8 @@ pub use host::generated::golem::core::types as wire; pub use decode::{ DecodeError, GraphDecoder, decode_graph, decode_metadata, decode_typed, decode_value, }; +#[cfg(all(feature = "guest", not(feature = "host")))] +pub use encode::encode_value_async; pub use encode::{ EncodeError, GraphEncoder, encode_graph, encode_metadata, encode_typed, encode_value, }; @@ -57,16 +59,25 @@ pub use decode::{ reject_quota_handles_in_value_tree, reject_secret_handles_in_value_tree, }; #[cfg(all(feature = "host", not(feature = "guest")))] -pub use encode::encode_value_with; +pub use encode::{encode_value_with, encode_value_with_streams}; #[cfg(all(feature = "host", not(feature = "guest")))] pub use host_support::{ - QuotaTokenHandleDropper, QuotaTokenHandleRep, QuotaTokenResolver, SecretHandleDropper, - SecretHandleRep, SecretResolver, + QuotaTokenHandleDropper, QuotaTokenHandleRep, QuotaTokenResolver, SchemaValueStreamResolver, + SecretHandleDropper, SecretHandleRep, SecretResolver, }; #[cfg(all(feature = "guest", not(feature = "host")))] pub use guest_support::{GuestQuotaTokenHandle, GuestSecretHandle}; +/// Create a native Component Model stream carrying schema value trees. +#[cfg(all(feature = "guest", not(feature = "host")))] +pub fn new_schema_value_stream() -> ( + wit_bindgen::StreamWriter, + wit_bindgen::StreamReader, +) { + guest::generated::wit_stream::new() +} + /// Host-side bridge for the opaque `golem:core/types.quota-token` resource. /// /// The resource is generated by this crate's host bindings and mapped (via the @@ -78,8 +89,25 @@ pub use guest_support::{GuestQuotaTokenHandle, GuestSecretHandle}; #[cfg(all(feature = "host", not(feature = "guest")))] mod host_support { use crate::schema::schema_value::{QuotaTokenValuePayload, SecretValuePayload}; + use crate::schema::{SchemaValueStream, SchemaValueStreamHandleRep}; use wasmtime::component::Resource; + pub trait SchemaValueStreamResolver { + type Error: std::fmt::Display; + + fn handle_from_stream( + &mut self, + stream: SchemaValueStream, + ) -> Result, Self::Error>; + + fn stream_from_handle( + &mut self, + handle: Resource, + ) -> Result; + + fn drop_stream_handle(&mut self, handle: Resource); + } + /// Opaque host representation backing the `quota-token` WIT resource. /// /// The boxed payload is owned and interpreted solely by the embedder; this @@ -264,12 +292,11 @@ mod guest_support { use crate::schema::metadata::TypeId; use crate::schema::schema_type::{QuotaTokenSpec, SchemaType}; use crate::schema::schema_value::SchemaValue; - use std::cell::RefCell; - use std::rc::Rc; + use std::sync::{Arc, Mutex}; #[derive(Clone)] pub struct GuestQuotaTokenHandle { - inner: Rc>>, + inner: Arc>>, } // The methods below expose the affine owned handle so the codec (this @@ -287,7 +314,7 @@ mod guest_support { #[doc(hidden)] pub fn new(handle: wire::QuotaToken) -> Self { Self { - inner: Rc::new(RefCell::new(Some(handle))), + inner: Arc::new(Mutex::new(Some(handle))), } } @@ -295,13 +322,19 @@ mod guest_support { /// already transferred (consumed) by a previous encode. #[doc(hidden)] pub fn take(&self) -> Option { - self.inner.borrow_mut().take() + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .take() } /// Whether the handle is still present (not yet transferred). #[doc(hidden)] pub fn is_present(&self) -> bool { - self.inner.borrow().is_some() + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .is_some() } /// Run `f` with a shared reference to the owned handle, if it is still @@ -313,14 +346,18 @@ mod guest_support { /// ownership of it. #[doc(hidden)] pub fn with_handle(&self, f: impl FnOnce(&wire::QuotaToken) -> R) -> Option { - self.inner.borrow().as_ref().map(f) + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .as_ref() + .map(f) } /// Identity of the shared cell, used to detect the same token appearing /// more than once in a single value tree. #[doc(hidden)] pub fn cell_id(&self) -> *const () { - Rc::as_ptr(&self.inner).cast() + Arc::as_ptr(&self.inner).cast() } } @@ -337,7 +374,7 @@ mod guest_support { impl PartialEq for GuestQuotaTokenHandle { fn eq(&self, other: &Self) -> bool { - Rc::ptr_eq(&self.inner, &other.inner) + Arc::ptr_eq(&self.inner, &other.inner) } } @@ -386,7 +423,7 @@ mod guest_support { #[derive(Clone)] pub struct GuestSecretHandle { - inner: Rc>>, + inner: Arc>>, } impl GuestSecretHandle { @@ -394,7 +431,7 @@ mod guest_support { #[doc(hidden)] pub fn new(handle: wire::Secret) -> Self { Self { - inner: Rc::new(RefCell::new(Some(handle))), + inner: Arc::new(Mutex::new(Some(handle))), } } @@ -402,27 +439,37 @@ mod guest_support { /// already transferred by a previous encode. #[doc(hidden)] pub fn take(&self) -> Option { - self.inner.borrow_mut().take() + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .take() } /// Whether the handle is still present (not yet transferred). #[doc(hidden)] pub fn is_present(&self) -> bool { - self.inner.borrow().is_some() + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .is_some() } /// Run `f` with a shared reference to the owned handle, if it is still /// present. #[doc(hidden)] pub fn with_handle(&self, f: impl FnOnce(&wire::Secret) -> R) -> Option { - self.inner.borrow().as_ref().map(f) + self.inner + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .as_ref() + .map(f) } /// Identity of the shared cell, used to detect the same secret /// appearing more than once in a single value tree. #[doc(hidden)] pub fn cell_id(&self) -> *const () { - Rc::as_ptr(&self.inner).cast() + Arc::as_ptr(&self.inner).cast() } } @@ -439,7 +486,7 @@ mod guest_support { impl PartialEq for GuestSecretHandle { fn eq(&self, other: &Self) -> bool { - Rc::ptr_eq(&self.inner, &other.inner) + Arc::ptr_eq(&self.inner, &other.inner) } } diff --git a/golem-schema/wit/deps/golem-core-v2/golem-core-v2.wit b/golem-schema/wit/deps/golem-core-v2/golem-core-v2.wit index f393d304d6..2b731300f4 100644 --- a/golem-schema/wit/deps/golem-core-v2/golem-core-v2.wit +++ b/golem-schema/wit/deps/golem-core-v2/golem-core-v2.wit @@ -125,6 +125,17 @@ interface types { /// and reveal it only through capability-gated host interfaces. resource secret; + /// An affine wrapper around a native Component Model stream of schema + /// values. The indirection lets a stream occur anywhere in a recursive + /// `schema-value-tree` while preserving the native stream endpoint. + resource schema-value-stream { + /// Wraps any native schema-value reader for placement in a value tree. + wrap: static async func(reader: stream) -> own; + + /// Consumes a schema-value-stream wrapper and returns its native reader. + unwrap: static async func(value: own) -> stream; + } + // ============================================================ // Schema graph (self-contained type carrier) // ============================================================ @@ -549,8 +560,7 @@ interface types { // Capability nodes secret-value(own), quota-token-handle(own), - - // WASI P3 stubs (parseable only in the schema; no constructible values). + stream-value(own), } record variant-value-payload { @@ -609,4 +619,5 @@ interface types { graph: schema-graph, value: schema-value-tree, } + } diff --git a/golem-service-base/src/grpc/client.rs b/golem-service-base/src/grpc/client.rs index 28528acd46..5cbec46ef8 100644 --- a/golem-service-base/src/grpc/client.rs +++ b/golem-service-base/src/grpc/client.rs @@ -202,6 +202,35 @@ impl MultiTargetGrpcClient { endpoint: Uri, f: F, ) -> Result + where + F: for<'a> Fn(&'a mut T) -> Pin> + 'a + Send>> + + Send, + { + self.call_with_retry(description, endpoint, true, f).await + } + + /// Performs one attempt, even for `Unavailable`, for operations whose request stream cannot + /// be replayed. + pub async fn call_without_retry( + &self, + description: impl AsRef, + endpoint: Uri, + f: F, + ) -> Result + where + F: for<'a> Fn(&'a mut T) -> Pin> + 'a + Send>> + + Send, + { + self.call_with_retry(description, endpoint, false, f).await + } + + async fn call_with_retry( + &self, + description: impl AsRef, + endpoint: Uri, + retry_on_unavailable: bool, + f: F, + ) -> Result where F: for<'a> Fn(&'a mut T) -> Pin> + 'a + Send>> + Send, @@ -232,7 +261,7 @@ impl MultiTargetGrpcClient { Err(e) => { if requires_reconnect(&e) { self.clients.remove_async(&endpoint).await; - if !retries.failed_attempt().await { + if !retry_on_unavailable || !retries.failed_attempt().await { span.in_scope(|| { warn!("gRPC call failed: {:?}, no more retries", e); }); @@ -435,3 +464,37 @@ impl EnabledGrpcClientTlsConfig { fn requires_reconnect(e: &Status) -> bool { e.code() == Code::Unavailable } + +#[cfg(test)] +mod tests { + use super::{GrpcClientConfig, MultiTargetGrpcClient}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use test_r::test; + use tonic::Status; + + #[derive(Clone)] + struct TestClient; + + #[test] + async fn multi_target_call_without_retry_attempts_unavailable_request_once() { + let client = + MultiTargetGrpcClient::new("test", |_, _| TestClient, GrpcClientConfig::default()); + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_call = attempts.clone(); + + let result = client + .call_without_retry( + "non-replayable", + "http://127.0.0.1:1".parse().unwrap(), + move |_| { + attempts_for_call.fetch_add(1, Ordering::Relaxed); + Box::pin(async { Err::<(), _>(Status::unavailable("failed")) }) + }, + ) + .await; + + assert!(result.is_err()); + assert_eq!(attempts.load(Ordering::Relaxed), 1); + } +} diff --git a/golem-service-base/src/model/agent_secret.rs b/golem-service-base/src/model/agent_secret.rs index 68122e536e..589ec95c71 100644 --- a/golem-service-base/src/model/agent_secret.rs +++ b/golem-service-base/src/model/agent_secret.rs @@ -51,16 +51,18 @@ impl From for golem_common::model::agent_secret::AgentSecretDto { } } -impl From for golem_api_grpc::proto::golem::registry::AgentSecret { - fn from(value: AgentSecret) -> Self { - Self { +impl TryFrom for golem_api_grpc::proto::golem::registry::AgentSecret { + type Error = String; + + fn try_from(value: AgentSecret) -> Result { + Ok(Self { agent_secret_id: Some(value.id.into()), environment_id: Some(value.environment_id.into()), path: value.path.0, revision: value.revision.into(), secret_type: Some(value.secret_type.into()), - secret_value: value.secret_value.map(Into::into), - } + secret_value: value.secret_value.map(TryInto::try_into).transpose()?, + }) } } diff --git a/golem-service-base/src/model/component.rs b/golem-service-base/src/model/component.rs index 43363b1fd5..dfcc78f079 100644 --- a/golem-service-base/src/model/component.rs +++ b/golem-service-base/src/model/component.rs @@ -135,14 +135,16 @@ impl TryFrom for Component { } } -impl From for golem_api_grpc::proto::golem::component::Component { - fn from(value: Component) -> Self { - Self { +impl TryFrom for golem_api_grpc::proto::golem::component::Component { + type Error = String; + + fn try_from(value: Component) -> Result { + Ok(Self { component_id: Some(value.id.into()), revision: value.revision.into(), component_name: value.component_name.0, component_size: value.component_size, - metadata: Some(value.metadata.into()), + metadata: Some(value.metadata.try_into()?), account_id: Some(value.account_id.into()), account_email: value.account_email.into_inner(), application_id: Some(value.application_id.into()), @@ -155,6 +157,6 @@ impl From for golem_api_grpc::proto::golem::component::Component { wasm_hash: Some(value.wasm_hash.into()), hash: Some(value.hash.into()), object_store_key: value.object_store_key, - } + }) } } diff --git a/golem-service-base/src/model/environment.rs b/golem-service-base/src/model/environment.rs index 24e2bd5f08..0608b0062e 100644 --- a/golem-service-base/src/model/environment.rs +++ b/golem-service-base/src/model/environment.rs @@ -30,18 +30,24 @@ pub struct EnvironmentState { pub tool_deployment: Option, } -impl From for golem_api_grpc::proto::golem::registry::EnvironmentState { - fn from(value: EnvironmentState) -> Self { - Self { +impl TryFrom for golem_api_grpc::proto::golem::registry::EnvironmentState { + type Error = String; + + fn try_from(value: EnvironmentState) -> Result { + Ok(Self { agent_deployment_details: value .agent_deployment_details .into_values() .map(Into::into) .collect(), - agent_secrets: value.agent_secrets.into_values().map(Into::into).collect(), + agent_secrets: value + .agent_secrets + .into_values() + .map(TryInto::try_into) + .collect::>()?, retry_policies: value.retry_policies.into_iter().map(Into::into).collect(), tool_deployment: value.tool_deployment.map(Into::into), - } + }) } } @@ -97,7 +103,8 @@ mod tests { tool_deployment: Some(tool_deployment.clone()), }; - let proto: golem_api_grpc::proto::golem::registry::EnvironmentState = state.into(); + let proto: golem_api_grpc::proto::golem::registry::EnvironmentState = + state.try_into().unwrap(); let decoded = EnvironmentState::try_from(proto).unwrap(); assert_eq!(decoded.tool_deployment, Some(tool_deployment)); diff --git a/golem-skills/skills/moonbit/golem-invoke-agent-moonbit/SKILL.md b/golem-skills/skills/moonbit/golem-invoke-agent-moonbit/SKILL.md index 1719e624ad..1de8da540a 100644 --- a/golem-skills/skills/moonbit/golem-invoke-agent-moonbit/SKILL.md +++ b/golem-skills/skills/moonbit/golem-invoke-agent-moonbit/SKILL.md @@ -19,7 +19,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using MoonBit syntax. Multiple return values are rendered as a MoonBit tuple, for example `(1, "ok")`. Methods returning `Unit` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `Unit` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one MoonBit value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consume_values - +cat input.bin | golem agent invoke 'MyAgent()' consume_bytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as MoonBit values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produce_values --stdout-format value +golem agent invoke 'MyAgent()' produce_bytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -78,6 +96,8 @@ golem agent invoke 'staging/MyAgent("user-123")' get_status | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/golem-skills/skills/rust/golem-invoke-agent-rust/SKILL.md b/golem-skills/skills/rust/golem-invoke-agent-rust/SKILL.md index d811a2a1e2..52e4f2ad47 100644 --- a/golem-skills/skills/rust/golem-invoke-agent-rust/SKILL.md +++ b/golem-skills/skills/rust/golem-invoke-agent-rust/SKILL.md @@ -19,7 +19,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using Rust syntax. Multiple return values are rendered as a Rust tuple, for example `(1, "ok")`. Methods returning `()` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `()` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one Rust value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consume_values - +cat input.bin | golem agent invoke 'MyAgent()' consume_bytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as Rust values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produce_values --stdout-format value +golem agent invoke 'MyAgent()' produce_bytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -73,6 +91,8 @@ golem agent invoke 'staging/MyAgent("user-123")' get_status | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/golem-skills/skills/scala/golem-invoke-agent-scala/SKILL.md b/golem-skills/skills/scala/golem-invoke-agent-scala/SKILL.md index 4fc9745094..8379c31dcc 100644 --- a/golem-skills/skills/scala/golem-invoke-agent-scala/SKILL.md +++ b/golem-skills/skills/scala/golem-invoke-agent-scala/SKILL.md @@ -19,7 +19,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using Scala syntax. Multiple return values are rendered as a Scala tuple, for example `(1, "ok")`. Methods returning `Unit` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `Unit` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one Scala value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consumeValues - +cat input.bin | golem agent invoke 'MyAgent()' consumeBytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as Scala values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produceValues --stdout-format value +golem agent invoke 'MyAgent()' produceBytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -73,6 +91,8 @@ golem agent invoke 'staging/MyAgent("user-123")' getStatus | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/golem-skills/skills/ts/golem-invoke-agent-ts/SKILL.md b/golem-skills/skills/ts/golem-invoke-agent-ts/SKILL.md index 96e1eca84b..42a575ad32 100644 --- a/golem-skills/skills/ts/golem-invoke-agent-ts/SKILL.md +++ b/golem-skills/skills/ts/golem-invoke-agent-ts/SKILL.md @@ -19,7 +19,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using TypeScript syntax. Multiple return values are rendered as a TypeScript tuple, for example `[1, "ok"]`. Methods returning `void` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `void` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one TypeScript value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf '{ value: 1 }\n{ value: 2 }\n' | golem agent invoke 'MyAgent()' consumeValues - +cat input.bin | golem agent invoke 'MyAgent()' consumeBytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as TypeScript values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produceValues --stdout-format value +golem agent invoke 'MyAgent()' produceBytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -73,6 +91,8 @@ golem agent invoke 'staging/MyAgent("user-123")' getStatus | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/golem-skills/tests/harness/src/executor.ts b/golem-skills/tests/harness/src/executor.ts index 7dfa844bcb..54afa62c6e 100644 --- a/golem-skills/tests/harness/src/executor.ts +++ b/golem-skills/tests/harness/src/executor.ts @@ -83,6 +83,22 @@ function parseJsonCommandOutput(output: string): T | undefined { } function extractInvokeJsonResult(output: string): unknown { + const lifecycleResult = output + .trim() + .split(/\r?\n/) + .map((line) => tryParseJson(line.trim())) + .find( + (document) => + document !== undefined && + typeof document === "object" && + document !== null && + (document as Record).$type === "agent.invoke-session" && + (document as Record).kind === "result", + ); + if (lifecycleResult && typeof lifecycleResult === "object") { + return (lifecycleResult as Record).value; + } + const parsed = parseJsonCommandOutput(output); if (!parsed || typeof parsed !== "object") { return parsed; diff --git a/golem-skills/tests/harness/tests/variables-integration.test.ts b/golem-skills/tests/harness/tests/variables-integration.test.ts index bc1ff6eefb..b66f9074fd 100644 --- a/golem-skills/tests/harness/tests/variables-integration.test.ts +++ b/golem-skills/tests/harness/tests/variables-integration.test.ts @@ -309,6 +309,62 @@ describe("Variable substitution integration", () => { assert.equal(result.status, "pass", result.stepResults[0]?.error); }); + it("extracts invoke_json results from invocation-session lifecycle documents", async () => { + const driver = new StubDriver(); + const watcher = new SkillWatcher(workspace); + const opts: ScenarioExecutorOptions = { agent: "amp", language: "ts" }; + const executor = createExecutor(driver, watcher, workspace, bootstrapSkillSourceDirs, opts); + + (executor as unknown as Record)["findGolemProjectDir"] = async () => workspace; + (executor as unknown as Record)["runLocalCommand"] = async () => ({ + success: true, + stdout: [ + { $type: "agent.invoke-session", kind: "accepted", idempotencyKey: "abc-123" }, + { + $type: "agent.invoke-session", + kind: "result", + idempotencyKey: "abc-123", + value: { count: 2, active: true }, + }, + { + $type: "agent.invoke-session", + kind: "finished", + idempotencyKey: "abc-123", + outcome: "success", + }, + ] + .map((document) => JSON.stringify(document)) + .join("\n"), + stderr: "", + output: "", + exitCode: 0, + }); + + const spec: ScenarioSpec = { + name: "invoke-json-session-output", + settings: { cleanup: false }, + steps: [ + { + id: "invoke-json", + tag: "invoke_json" as const, + invoke_json: { + agent: 'CounterAgent("test")', + method: "get", + }, + expect: { + result_json: [ + { path: "$.count", equals: 2 }, + { path: "$.active", equals: true }, + ], + }, + }, + ], + }; + + const result = await executor.execute(spec); + assert.equal(result.status, "pass", result.stepResults[0]?.error); + }); + it("resolves language-conditional invoke_json method names before execution", async () => { const driver = new StubDriver(); const watcher = new SkillWatcher(workspace); diff --git a/golem-worker-executor-test-utils/Cargo.toml b/golem-worker-executor-test-utils/Cargo.toml index ca823f317c..013a7063a6 100644 --- a/golem-worker-executor-test-utils/Cargo.toml +++ b/golem-worker-executor-test-utils/Cargo.toml @@ -35,6 +35,7 @@ serde = { workspace = true } serde_json = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true } +tokio-stream = { workspace = true } tokio-util = { workspace = true } tonic = { workspace = true } tonic-tracing-opentelemetry = { workspace = true } diff --git a/golem-worker-executor-test-utils/src/dsl_impl.rs b/golem-worker-executor-test-utils/src/dsl_impl.rs index 08c1e09f52..6c4eb77fa7 100644 --- a/golem-worker-executor-test-utils/src/dsl_impl.rs +++ b/golem-worker-executor-test-utils/src/dsl_impl.rs @@ -17,7 +17,12 @@ use crate::component_writer::CachedAnalysis; use anyhow::anyhow; use applying::Apply; use bytes::Bytes; -use golem_api_grpc::proto::golem::worker::{LogEvent, UpdateMode}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::worker::invocation_session_completion::Outcome; +use golem_api_grpc::proto::golem::worker::{ + InvocationRequest, InvocationSessionResult, InvocationStart, LogEvent, UpdateMode, + invocation_request, invocation_response, +}; use golem_api_grpc::proto::golem::workerexecutor; use golem_api_grpc::proto::golem::workerexecutor::v1::{ CancelInvocationRequest, CompletePromiseRequest, ConnectWorkerRequest, CreateWorkerRequest, @@ -60,6 +65,8 @@ use golem_test_framework::dsl::{AgentResult, TestDsl, WorkerLogEventStream}; use golem_test_framework::model::IFSEntry; use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; use tonic::Streaming; use uuid::Uuid; @@ -84,6 +91,84 @@ fn invocation_agent_id( .map_err(|err| anyhow!("Invalid agent id: {err}")) } +impl TestWorkerExecutor { + pub async fn invoke_agent_session( + &self, + start: InvocationStart, + ) -> anyhow::Result { + let (requests, receiver) = mpsc::channel(1); + let request = InvocationRequest { + request: Some(invocation_request::Request::Start(start)), + }; + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&request) + .map_err(anyhow::Error::msg)?; + requests + .send(request) + .await + .map_err(|_| anyhow!("invocation session request ended before start"))?; + let mut responses = self + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let mut result = None; + let mut terminal = None; + + while let Some(response) = responses.message().await? { + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + match response.response { + Some(invocation_response::Response::Accepted(_)) => {} + Some(invocation_response::Response::Rejected(rejected)) => { + terminal = Some(Err(anyhow!( + "Agent invocation rejected: {}", + rejected.error + ))); + } + Some(invocation_response::Response::Result(invocation_result)) => { + result = Some(invocation_result); + } + Some(invocation_response::Response::Finished(finished)) => { + terminal = Some(match finished.outcome { + Some(Outcome::Success(_)) => result + .take() + .ok_or_else(|| anyhow!("invocation completed without a result")), + Some(Outcome::Failure(failure)) => { + Err(anyhow!("Agent invocation failed: {failure:?}")) + } + None => Err(anyhow!("invocation completion has no outcome")), + }); + } + Some( + invocation_response::Response::OutputItem(_) + | invocation_response::Response::OutputEnd(_) + | invocation_response::Response::OutputError(_) + | invocation_response::Response::InputAck(_) + | invocation_response::Response::StreamCancel(_), + ) => { + return Err(anyhow!( + "non-streaming test invocation received a stream frame" + )); + } + Some(invocation_response::Response::AttachmentRevoked(_)) => { + unreachable!("response validation rejects attachment revocation") + } + None => unreachable!("response validation rejects empty frames"), + } + } + + terminal.unwrap_or_else(|| { + Err(anyhow!( + "invocation session response ended before completion" + )) + }) + } +} + #[async_trait::async_trait] impl TestDsl for TestWorkerExecutor { type WorkerError = WorkerExecutorError; @@ -445,40 +530,27 @@ impl TestDsl for TestWorkerExecutor { let (_graph, value) = params.into_parts(); let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - value.into(); - - let result = self - .client - .clone() - .invoke_agent(workerexecutor::v1::InvokeAgentRequest { - agent_id: Some(agent_id.clone().into()), - method_name: Some(method_name.to_string()), - method_parameters: Some(proto_method_parameters), - mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32, - schedule_at: None, - idempotency_key: Some(idempotency_key.clone().into()), - component_owner_account_id: Some(component.account_id.into()), - environment_id: Some(component.environment_id.into()), - auth_ctx: Some(self.auth_ctx().into()), - context: None, - principal: None, - freshness_disposition: workerexecutor::v1::InvocationFreshnessDisposition::MayExist + value.try_into().map_err(anyhow::Error::msg)?; + + self.invoke_agent_session(InvocationStart { + agent_id: Some(agent_id.clone().into()), + method_name: Some(method_name.to_string()), + input: Some(proto_method_parameters), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32, + schedule_at: None, + idempotency_key: Some(idempotency_key.clone().into()), + component_owner_account_id: Some(component.account_id.into()), + environment_id: Some(component.environment_id.into()), + auth_ctx: Some(self.auth_ctx().into()), + context: None, + principal: None, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist as i32, - config: Vec::new(), - }) - .await; - - let result = result?.into_inner(); - - match result.result { - None => Err(anyhow!( - "No response from golem-worker-executor invoke_agent call" - )), - Some(workerexecutor::v1::invoke_agent_response::Result::Success(_)) => Ok(()), - Some(workerexecutor::v1::invoke_agent_response::Result::Failure(error)) => { - Err(anyhow!("Failed converting error: {error:?}")) - } - } + config: Vec::new(), + }) + .await?; + Ok(()) } #[tracing::instrument(level = "info", skip_all, fields(component_id = %component.id, %agent_id, method_name))] @@ -499,15 +571,13 @@ impl TestDsl for TestWorkerExecutor { let (_graph, value) = params.into_parts(); let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - value.into(); + value.try_into().map_err(anyhow::Error::msg)?; let result = self - .client - .clone() - .invoke_agent(workerexecutor::v1::InvokeAgentRequest { + .invoke_agent_session(InvocationStart { agent_id: Some(worker_agent_id.clone().into()), method_name: Some(method_name.to_string()), - method_parameters: Some(proto_method_parameters), + input: Some(proto_method_parameters), mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, schedule_at: None, idempotency_key: Some(key.into()), @@ -516,32 +586,24 @@ impl TestDsl for TestWorkerExecutor { auth_ctx: Some(self.auth_ctx().into()), context: None, principal: principal.map(Into::into), - freshness_disposition: workerexecutor::v1::InvocationFreshnessDisposition::MayExist - as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, config: Vec::new(), }) - .await; - - let result = result?.into_inner(); + .await?; - match result.result { - None => Err(anyhow!( - "No response from golem-worker-executor invoke_agent call" - )), - Some(workerexecutor::v1::invoke_agent_response::Result::Success(success)) => { - let value = match success.result { - Some(proto_val) => Some( - SchemaValue::try_from(proto_val) - .map_err(|err| anyhow!("SchemaValue conversion error: {err}"))?, - ), - None => None, - }; - Ok(AgentResult::new(value)) - } - Some(workerexecutor::v1::invoke_agent_response::Result::Failure(error)) => { - Err(anyhow!("Agent invocation failed: {error:?}")) - } - } + let value = match result.result { + Some( + golem_api_grpc::proto::golem::worker::invocation_session_result::Result::MethodResult(proto_value), + ) => Some( + SchemaValue::try_from(proto_value) + .map_err(|err| anyhow!("SchemaValue conversion error: {err}"))?, + ), + Some(golem_api_grpc::proto::golem::worker::invocation_session_result::Result::NoResult(_)) + | None => None, + }; + Ok(AgentResult::new(value)) } #[tracing::instrument(level = "info", skip_all, fields(%agent_id))] diff --git a/golem-worker-executor-test-utils/src/lib.rs b/golem-worker-executor-test-utils/src/lib.rs index 34b4d66273..75ad2f4b71 100644 --- a/golem-worker-executor-test-utils/src/lib.rs +++ b/golem-worker-executor-test-utils/src/lib.rs @@ -2163,7 +2163,10 @@ impl Bootstrap &mut linker, >::durable_ctx_mut, )?; - golem_schema::schema::wit::wire::add_to_linker::<_, HasSelf>>( + golem_schema::schema::wit::wire::add_to_linker::< + _, + golem_worker_executor::durable_host::CoreTypesHost, + >( &mut linker, >::durable_ctx_mut, )?; diff --git a/golem-worker-executor/Cargo.toml b/golem-worker-executor/Cargo.toml index 42e71c2eb4..4d9f73a471 100644 --- a/golem-worker-executor/Cargo.toml +++ b/golem-worker-executor/Cargo.toml @@ -36,6 +36,7 @@ golem-service-base = { workspace = true, features = ["worker-executor"] } anyhow = { workspace = true } applying = { workspace = true } arc-swap = { workspace = true } +async-broadcast = { workspace = true } async-lock = { workspace = true } async-recursion = { workspace = true } async-scoped = { workspace = true, features = ["use-tokio"] } diff --git a/golem-worker-executor/config/worker-executor.sample.env b/golem-worker-executor/config/worker-executor.sample.env index 85aee455f5..3f65cc49a6 100644 --- a/golem-worker-executor/config/worker-executor.sample.env +++ b/golem-worker-executor/config/worker-executor.sample.env @@ -84,6 +84,7 @@ GOLEM__LIMITS__EVENT_BROADCAST_CAPACITY=1024 GOLEM__LIMITS__EVENT_HISTORY_SIZE=128 GOLEM__LIMITS__FUEL_TO_BORROW=10000 GOLEM__LIMITS__INVOCATION_RESULT_BROADCAST_CAPACITY=100000 +GOLEM__LIMITS__LIVE_STREAM_EVENT_BROADCAST_CAPACITY=32 GOLEM__LIMITS__MAX_ACTIVE_WORKERS=1024 GOLEM__LIMITS__MAX_CONCURRENT_STREAMS=1024 GOLEM__LIMITS__MAX_INVOCATION_CONTEXT_STACK_DEPTH=1024 @@ -329,6 +330,7 @@ GOLEM__LIMITS__EVENT_BROADCAST_CAPACITY=1024 GOLEM__LIMITS__EVENT_HISTORY_SIZE=128 GOLEM__LIMITS__FUEL_TO_BORROW=10000 GOLEM__LIMITS__INVOCATION_RESULT_BROADCAST_CAPACITY=100000 +GOLEM__LIMITS__LIVE_STREAM_EVENT_BROADCAST_CAPACITY=32 GOLEM__LIMITS__MAX_ACTIVE_WORKERS=1024 GOLEM__LIMITS__MAX_CONCURRENT_STREAMS=1024 GOLEM__LIMITS__MAX_INVOCATION_CONTEXT_STACK_DEPTH=1024 @@ -544,6 +546,7 @@ GOLEM__LIMITS__EVENT_BROADCAST_CAPACITY=1024 GOLEM__LIMITS__EVENT_HISTORY_SIZE=128 GOLEM__LIMITS__FUEL_TO_BORROW=10000 GOLEM__LIMITS__INVOCATION_RESULT_BROADCAST_CAPACITY=100000 +GOLEM__LIMITS__LIVE_STREAM_EVENT_BROADCAST_CAPACITY=32 GOLEM__LIMITS__MAX_ACTIVE_WORKERS=1024 GOLEM__LIMITS__MAX_CONCURRENT_STREAMS=1024 GOLEM__LIMITS__MAX_INVOCATION_CONTEXT_STACK_DEPTH=1024 diff --git a/golem-worker-executor/config/worker-executor.toml b/golem-worker-executor/config/worker-executor.toml index 2d58be938b..6c2fdf8d17 100644 --- a/golem-worker-executor/config/worker-executor.toml +++ b/golem-worker-executor/config/worker-executor.toml @@ -143,6 +143,7 @@ event_broadcast_capacity = 1024 event_history_size = 128 fuel_to_borrow = 10000 invocation_result_broadcast_capacity = 100000 +live_stream_event_broadcast_capacity = 32 max_active_workers = 1024 max_concurrent_streams = 1024 max_invocation_context_stack_depth = 1024 @@ -509,6 +510,7 @@ without_time = false # event_history_size = 128 # fuel_to_borrow = 10000 # invocation_result_broadcast_capacity = 100000 +# live_stream_event_broadcast_capacity = 32 # max_active_workers = 1024 # max_concurrent_streams = 1024 # max_invocation_context_stack_depth = 1024 @@ -845,6 +847,7 @@ without_time = false # event_history_size = 128 # fuel_to_borrow = 10000 # invocation_result_broadcast_capacity = 100000 +# live_stream_event_broadcast_capacity = 32 # max_active_workers = 1024 # max_concurrent_streams = 1024 # max_invocation_context_stack_depth = 1024 diff --git a/golem-worker-executor/src/durable_host/concurrent/call.rs b/golem-worker-executor/src/durable_host/concurrent/call.rs index dd34686a63..2a69804c48 100644 --- a/golem-worker-executor/src/durable_host/concurrent/call.rs +++ b/golem-worker-executor/src/durable_host/concurrent/call.rs @@ -118,14 +118,14 @@ pub struct CallHandle { /// pre-call index and `end_durable_function` only uses it to commit at the right boundary. pub(super) begin_index: OplogIndex, pub(super) is_live: bool, - /// `true` when a `Start` entry was actually appended. It is `false` while snapshotting (where - /// nothing is persisted) and for replay handles. + /// `true` when a `Start` entry was actually appended. It is `false` during unpersisted + /// execution and for replay handles. pub(super) persisted: bool, /// Tracks the (possibly deferred) blob upload of this call's request payload, started when the /// `Start` was reserved. Awaited before the matching `End` / `Cancelled` is appended so an /// upload failure surfaces at the call site rather than only at the leaf oplog's commit barrier. - /// `PendingUpload::already_durable()` (a no-op) for replay handles, snapshotting, and inline - /// requests. + /// `PendingUpload::already_durable()` (a no-op) for replay handles, unpersisted execution, and + /// inline requests. pub(super) request_upload: PendingUpload, /// Replay-side resolver receiver; `Some` only for replay handles. pub(super) replay: Option, @@ -261,7 +261,7 @@ impl CallExecutionScope { /// Builds an *unregistered* atomic-region lease: it preserves the call's initiation-time region /// for trap/retry classification (matching the immutable capture used before leases existed) but /// is not a member of any region registry, so it never transfers or detaches on region close. -/// Used for replay and snapshotting handles, which do not participate in the live in-flight +/// Used for replay and unpersisted handles, which do not participate in the live in-flight /// member guard. pub(super) fn unregistered_atomic_lease( atomic_region: Option, @@ -277,7 +277,7 @@ pub(super) fn unregistered_atomic_lease( struct PreparedAccessStart { is_live: bool, - snapshotting: bool, + unpersisted: bool, oplog: Arc, public_state: PublicDurableWorkerState, replay_state: crate::durable_host::replay_state::ReplayState, @@ -770,7 +770,7 @@ impl CallHandle { /// ordering and settlement accounting survive. /// /// This handoff is only performed once *this* call is a persisted live barrier (never for - /// replay or snapshotting handles): until then the prior observer must stay armed so a guest + /// replay or unpersisted handles): until then the prior observer must stay armed so a guest /// cancellation landing before this call's `Start` still records its `CompletionDiscarded` /// marker and parks replay at the prior call. /// @@ -849,13 +849,14 @@ impl CallHandle { observational_owner, }; let is_live = durable_execution_state.is_live; - let snapshotting = durable_execution_state.snapshotting_mode; + let unpersisted = durable_execution_state.snapshotting_mode + || durable_execution_state.is_unpersisted_execution; let retry = InFunctionRetryController::new(function_type, durable_execution_state, Pair::FQFN); // A live persisted call initiated inside an open atomic region joins the region's member // registry: its lease starts owned by that region and follows the region's close // transitions (transfer to the enclosing region, or detachment at the outermost close). - let atomic_lease = if is_live && !snapshotting { + let atomic_lease = if is_live && !unpersisted { match atomic_region { Some(begin_index) => Some( ctx.state @@ -878,7 +879,7 @@ impl CallHandle { let live_host_calls = ctx.state.live_host_call_counter(); Ok(PreparedAccessStart { is_live, - snapshotting, + unpersisted, oplog: ctx.state.oplog.clone(), public_state: ctx.public_state.clone(), replay_state: ctx.state.replay_state.clone(), @@ -899,7 +900,7 @@ impl CallHandle { }) } - /// Snapshotting is the only state in which a live host call is intentionally not persisted. + /// Unpersisted executions take the live `persisted: false` branch. async fn execute_access_start( mut prepared: PreparedAccessStart, build_request: F, @@ -911,7 +912,7 @@ impl CallHandle { let starts_scope = opens_accessor_scope( prepared.retry.function_type(), prepared.retry.durable_execution_state().assume_idempotence, - prepared.snapshotting, + prepared.unpersisted, ); let scope_start = if starts_scope { Some(Self::execute_access_scope_start(&prepared).await?) @@ -931,6 +932,7 @@ impl CallHandle { DurableExecutionState { is_live: true, snapshotting_mode: previous.snapshotting_mode, + is_unpersisted_execution: previous.is_unpersisted_execution, assume_idempotence: previous.assume_idempotence, max_in_function_retry_delay: previous.max_in_function_retry_delay, }, @@ -973,7 +975,7 @@ impl CallHandle { })?; if is_live { - if prepared.snapshotting { + if prepared.unpersisted { let start_idx = prepared.oplog.current_oplog_index().await; let atomic_lease = unregistered_atomic_lease( execution_scope.atomic_region, @@ -1632,6 +1634,10 @@ impl CallHandle { .max_in_function_retry_delay, current_retry_policy_state, retry_properties: properties.clone(), + is_unpersisted_execution: self + .retry + .durable_execution_state() + .is_unpersisted_execution, worker, }; @@ -2708,25 +2714,37 @@ where D: HasData + ?Sized, Ctx: WorkerCtx, { - let (opens_scope, is_live, replay_handle, replay_state, oplog, public_state) = - store.with(|mut access| { - let ctx = get_ctx(access.data_mut()); - let opens_scope = ctx.state.opens_durable_scope(&function_type); - let is_live = ctx.state.is_live(); - let replay_handle = if opens_scope && !is_live { - ctx.state.take_durable_scope_replay_handle(begin_index) - } else { - None - }; - ( - opens_scope, - is_live, - replay_handle, - ctx.state.replay_state.clone(), - ctx.state.oplog.clone(), - ctx.public_state.clone(), - ) - }); + let ( + opens_scope, + is_live, + is_unpersisted_execution, + replay_handle, + replay_state, + oplog, + public_state, + ) = store.with(|mut access| { + let ctx = get_ctx(access.data_mut()); + let opens_scope = ctx.state.opens_durable_scope(&function_type); + let is_live = ctx.state.is_live(); + let replay_handle = if opens_scope && !is_live { + ctx.state.take_durable_scope_replay_handle(begin_index) + } else { + None + }; + ( + opens_scope, + is_live, + ctx.state.durability_is_suppressed(), + replay_handle, + ctx.state.replay_state.clone(), + ctx.state.oplog.clone(), + ctx.public_state.clone(), + ) + }); + + if is_unpersisted_execution { + return Ok(()); + } if opens_scope { if is_live { @@ -2873,21 +2891,24 @@ where D: HasData + ?Sized, Ctx: WorkerCtx, { - let (is_live, worker, replay_state) = store.with(|mut access| { + let (is_live, durability_is_suppressed, worker, replay_state) = store.with(|mut access| { let ctx = get_ctx(access.data_mut()); ( ctx.state.is_live(), + ctx.state.durability_is_suppressed(), ctx.public_state.worker(), ctx.state.replay_state.clone(), ) }); - if is_live { - worker - .add_to_oplog(OplogEntry::finish_span(span_id.clone())) - .await; - } else { - crate::get_oplog_entry_owned!(replay_state, OplogEntry::FinishSpan)?; + if !durability_is_suppressed { + if is_live { + worker + .add_to_oplog(OplogEntry::finish_span(span_id.clone())) + .await; + } else { + crate::get_oplog_entry_owned!(replay_state, OplogEntry::FinishSpan)?; + } } store.with(|mut access| { @@ -2909,9 +2930,9 @@ fn is_accessor_supported_function_type(function_type: &DurableFunctionType) -> b fn opens_accessor_scope( function_type: &DurableFunctionType, assume_idempotence: bool, - snapshotting: bool, + unpersisted: bool, ) -> bool { - !snapshotting + !unpersisted && ((*function_type == DurableFunctionType::WriteRemote && !assume_idempotence) || matches!(function_type, DurableFunctionType::WriteRemoteBatched(None))) } @@ -3364,20 +3385,22 @@ impl BegunCall { } /// Second phase on the live path: upload the request and append the eager host-call `Start` - /// (or, while snapshotting, persist nothing). + /// (or, while durability is suppressed, persist nothing). pub async fn start_live( self, ctx: &mut DurableWorkerCtx, request: Pair::Req, ) -> Result, WorkerExecutorError> { debug_assert!(self.is_live(), "start_live() called on a replay handle"); - let snapshotting = self.retry.durable_execution_state().snapshotting_mode; + let durable_execution_state = self.retry.durable_execution_state(); + let durability_is_suppressed = durable_execution_state.snapshotting_mode + || durable_execution_state.is_unpersisted_execution; // The host-call `Start` nests inside the enclosing durable scope captured at initiation // (its own opened scope, or the scope encoded in the function type), derived explicitly — // never from the set of temporally-open sibling scopes. `None` for a top-level unscoped call. let parent_start_index = self.execution_scope.parent_start_index; - let (start_idx, persisted, request_upload) = if snapshotting { - // Snapshotting mode persists nothing. + let (start_idx, persisted, request_upload) = if durability_is_suppressed { + // Snapshot and unpersisted execution write no durable call records. let oplog = ctx.state.oplog.clone(); ( oplog.current_oplog_index().await, @@ -3431,7 +3454,7 @@ impl BegunCall { None } } else { - // Snapshotting persists nothing; keep the initiation-time region for trap/retry + // Unpersisted execution writes nothing; keep the initiation-time region for trap/retry // classification without joining the live in-flight member guard. unregistered_atomic_lease( self.execution_scope.atomic_region, @@ -3510,7 +3533,7 @@ impl Drop for CallHandle { self.drop_sink.as_ref(), ); } - // Not persisted (snapshotting): there is nothing on disk to reconcile. + // Not persisted: there is nothing on disk to reconcile. } else { if opens_replay_durable_scope( self.retry.function_type(), diff --git a/golem-worker-executor/src/durable_host/concurrent/delivery.rs b/golem-worker-executor/src/durable_host/concurrent/delivery.rs index b3899bdc92..c9696d060a 100644 --- a/golem-worker-executor/src/durable_host/concurrent/delivery.rs +++ b/golem-worker-executor/src/durable_host/concurrent/delivery.rs @@ -99,7 +99,7 @@ pub struct CompletionDelivery { pub(super) enum CompletionDeliveryState { /// Live, armed: the `End` is persisted and a torn/failed delivery must record a marker. Live(Box), - /// Live, but the call was not persisted (snapshotting): nothing to reconcile. + /// Live, but the call was not persisted: nothing to reconcile. Unarmed, /// Replay of a normally delivered completion. ReplayDelivered, @@ -231,7 +231,7 @@ impl CompletionDelivery { /// recorded, replay re-executes the host code past this `End` deterministically and /// re-consumes the response internally, so no marker is needed. /// - /// Non-live tokens (replay, unpersisted snapshotting calls) settle immediately; if the + /// Non-live tokens (replay and unpersisted calls) settle immediately; if the /// accessor has no guest-visible host subtask (e.g. a spawned background task), the token /// settles without a marker, matching the pre-observer behavior of consuming it at the host /// return. diff --git a/golem-worker-executor/src/durable_host/concurrent/tests.rs b/golem-worker-executor/src/durable_host/concurrent/tests.rs index 26520a57f4..1f4a8fe570 100644 --- a/golem-worker-executor/src/durable_host/concurrent/tests.rs +++ b/golem-worker-executor/src/durable_host/concurrent/tests.rs @@ -28,6 +28,7 @@ fn durable_execution_state() -> DurableExecutionState { DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: true, max_in_function_retry_delay: Duration::from_secs(20), } @@ -172,6 +173,7 @@ fn live_unfinished_handle_with_atomic_region( let durable_execution_state = DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: false, max_in_function_retry_delay: Duration::ZERO, }; @@ -216,6 +218,7 @@ fn synthetic_finished_handle_with_scope( let durable_execution_state = DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: false, max_in_function_retry_delay: Duration::ZERO, }; @@ -1518,6 +1521,7 @@ fn can_reexecute_matches_internal_retry_eligibility() { DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence, max_in_function_retry_delay: Duration::ZERO, }, diff --git a/golem-worker-executor/src/durable_host/durability.rs b/golem-worker-executor/src/durable_host/durability.rs index fef76b9cfc..cf6f31c8e9 100644 --- a/golem-worker-executor/src/durable_host/durability.rs +++ b/golem-worker-executor/src/durable_host/durability.rs @@ -276,7 +276,7 @@ pub(crate) struct OpenCustomInvocationScope { #[derive(Debug)] pub struct LiveCustomDurableInvocation { scope_id: u64, - start_index: OplogIndex, + start_index: Option, } /// Classification of host function failures for semantic retry decisions @@ -605,6 +605,7 @@ pub enum AsyncRetryDecision { pub struct DurableExecutionState { pub is_live: bool, pub snapshotting_mode: bool, + pub is_unpersisted_execution: bool, /// Whether the executor assumes idempotence for remote writes. pub assume_idempotence: bool, /// Maximum delay for in-function retries. Delays exceeding this fall back to trap+replay. @@ -1308,7 +1309,7 @@ fn open_live_custom_durable_invocation( }; let resource = access.get().table().push(LiveCustomDurableInvocation { scope_id, - start_index, + start_index: Some(start_index), })?; access.get().state.custom_invocation_scopes.insert( scope_id, @@ -1338,12 +1339,15 @@ fn open_live_custom_durable_invocation( fn close_live_custom_durable_invocation( access: &mut Access<'_, U, HasSelf>>, resource: Resource, -) -> anyhow::Result { +) -> anyhow::Result> { + let resource_rep = resource.rep(); + let invocation = access.get().table().delete(resource)?; + let Some(start_index) = invocation.start_index else { + return Ok(None); + }; let current = access .as_context_mut() .guest_task_context::()?; - let resource_rep = resource.rep(); - let invocation = access.get().table().delete(resource)?; let registered = access .get() .state @@ -1355,9 +1359,7 @@ fn close_live_custom_durable_invocation( invocation.scope_id ) })?; - if registered.owner_start_index != invocation.start_index - || registered.resource_rep != resource_rep - { + if registered.owner_start_index != start_index || registered.resource_rep != resource_rep { return Err(anyhow::anyhow!( "custom invocation scope {} has stale attribution", invocation.scope_id @@ -1387,7 +1389,7 @@ fn close_live_custom_durable_invocation( .as_context_mut() .set_guest_task_context(Arc::new(next))?; } - Ok(invocation.start_index) + Ok(Some(start_index)) } impl durability::HostLiveCustomDurableInvocationWithStore @@ -1397,12 +1399,13 @@ impl durability::HostLiveCustomDurableInvocat mut access: Access<'_, U, Self>, resource: Resource, ) -> anyhow::Result<()> { - let start_index = close_live_custom_durable_invocation(&mut access, resource)?; - access - .get() - .state - .active_custom_invocations - .remove(&start_index); + if let Some(start_index) = close_live_custom_durable_invocation(&mut access, resource)? { + access + .get() + .state + .active_custom_invocations + .remove(&start_index); + } Ok(()) } @@ -1412,9 +1415,18 @@ impl durability::HostLiveCustomDurableInvocat response: golem_common::schema::wit::wire::TypedSchemaValue, forced_commit: bool, ) -> anyhow::Result<()> { - let (start_index, response, oplog, worker) = accessor.with(|mut access| { + let invocation = accessor.with(|mut access| { let start_index = close_live_custom_durable_invocation(&mut access, resource)?; let ctx = access.get(); + let Some(start_index) = start_index else { + golem_common::schema::wit::decode_typed_rejecting_quota_with(response, ctx) + .map_err(|e| { + anyhow::anyhow!( + "Failed to decode durable function response schema value: {e}" + ) + })?; + return Ok::<_, anyhow::Error>(None); + }; let invocation = ctx .state .active_custom_invocations @@ -1447,8 +1459,11 @@ impl durability::HostLiveCustomDurableInvocat })?; let oplog = ctx.state.oplog.clone(); let worker = ctx.public_state.worker(); - Ok::<_, anyhow::Error>((start_index, response, oplog, worker)) + Ok::<_, anyhow::Error>(Some((start_index, response, oplog, worker))) })?; + let Some((start_index, response, oplog, worker)) = invocation else { + return Ok(()); + }; concurrent::drain_dropped_call_events_access(accessor, accessor.getter()) .await @@ -1490,6 +1505,34 @@ impl durability::HostWithStore request: golem_common::schema::wit::wire::TypedSchemaValue, function_type: durability::DurableFunctionType, ) -> anyhow::Result { + let is_unpersisted = + accessor.with(|mut access| access.get().state.is_unpersisted_execution()); + if is_unpersisted { + let resource = accessor.with(|mut access| { + let ctx = access.get(); + golem_common::schema::wit::decode_typed_rejecting_quota_with(request, ctx) + .map_err(|e| { + anyhow::anyhow!( + "Failed to decode durable function request schema value: {e}" + ) + })?; + let function_type: DurableFunctionType = function_type.into(); + if is_write_side_effect(&function_type) { + DurableWorkerCtx::check_read_only_allows( + ctx, + "golem::durability::begin-custom-durable-invocation", + ) + .map_err(anyhow::Error::msg)?; + } + let resource = ctx.table().push(LiveCustomDurableInvocation { + scope_id: 0, + start_index: None, + })?; + Ok::<_, anyhow::Error>(resource) + })?; + return Ok(durability::CustomDurableInvocation::Live(resource)); + } + // Capture attribution before the first await. The immutable Arc is the initiation-time // snapshot inherited by this host task and its continuations. let custom_invocation_context = accessor.guest_task_context::()?; @@ -1799,9 +1842,11 @@ impl InFunctionRetryHost for DurableWorkerCtx { } fn durable_execution_state(&self) -> DurableExecutionState { + let is_unpersisted_execution = self.is_unpersisted_execution(); DurableExecutionState { - is_live: self.state.is_live() || self.state.snapshotting_mode, + is_live: self.state.is_live() || self.state.durability_is_suppressed(), snapshotting_mode: self.state.snapshotting_mode, + is_unpersisted_execution, assume_idempotence: self.state.assume_idempotence, max_in_function_retry_delay: self.state.config.max_in_function_retry_delay, } @@ -1823,7 +1868,7 @@ impl InFunctionRetryHost for DurableWorkerCtx { inside_atomic_region: bool, retry_policy_state: Option, ) { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { return; } @@ -1885,7 +1930,7 @@ impl DurabilityHost for DurableWorkerCtx { forced_commit: bool, ) -> Result<(), WorkerExecutorError> { self.end_function(function_type, begin_index).await?; - if !self.state.snapshotting_mode + if !self.state.durability_is_suppressed() && (function_type == &DurableFunctionType::WriteRemote || matches!(function_type, DurableFunctionType::WriteRemoteBatched(_)) || matches!( @@ -2029,6 +2074,10 @@ impl DurabilityHost for DurableWorkerCtx { let status = self.execution_status.read().unwrap(); status.create_await_interrupt_signal() }; + let live_stream_cancellation = self + .live_stream_tracker + .as_ref() + .map(|tracker| tracker.cancellation_token()); if self .state .invocation_deadline_exceeded @@ -2037,6 +2086,9 @@ impl DurabilityHost for DurableWorkerCtx { .state .tail_work_deadline_exceeded .load(std::sync::atomic::Ordering::Acquire) + || live_stream_cancellation + .as_ref() + .is_some_and(tokio_util::sync::CancellationToken::is_cancelled) { let status = self.execution_status.read().unwrap(); if matches!(&*status, ExecutionStatus::Interrupting { .. }) { @@ -2046,7 +2098,18 @@ impl DurabilityHost for DurableWorkerCtx { Timestamp::now_utc(), ))); } - interrupt_signal + match live_stream_cancellation { + Some(cancellation) => Box::pin(async move { + tokio::select! { + biased; + interrupt = interrupt_signal => interrupt, + () = cancellation.cancelled() => { + InterruptKind::Interrupt(Timestamp::now_utc()) + } + } + }), + None => interrupt_signal, + } } fn check_read_only_allows(&self, host_function: &str) -> Result<(), GolemSpecificWasmTrap> { @@ -2310,6 +2373,8 @@ pub struct TaskRetryContext { pub current_retry_policy_state: Option, /// Properties describing the error context (verb, URI, status code, etc.) for predicate evaluation. pub retry_properties: RetryProperties, + /// Whether retry bookkeeping belongs to an unpersisted execution and must not be recorded. + pub is_unpersisted_execution: bool, /// Reference to the worker that owns this task. pub worker: Arc>, } @@ -2351,6 +2416,7 @@ impl InFunctionRetryHost for TaskRetryContext { DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: true, max_in_function_retry_delay: self.max_in_function_retry_delay, } @@ -2362,6 +2428,11 @@ impl InFunctionRetryHost for TaskRetryContext { inside_atomic_region: bool, retry_policy_state: Option, ) { + if self.is_unpersisted_execution { + self.current_retry_policy_state = retry_policy_state; + return; + } + use golem_common::model::oplog::AgentError; let entry = OplogEntry::error( AgentError::TransientError("in-function retry".to_string()), @@ -2615,6 +2686,7 @@ mod tests { DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: self.assume_idempotence, max_in_function_retry_delay: self.max_in_function_retry_delay, } diff --git a/golem-worker-executor/src/durable_host/golem/retry_api.rs b/golem-worker-executor/src/durable_host/golem/retry_api.rs index 21ede011ad..f5270be020 100644 --- a/golem-worker-executor/src/durable_host/golem/retry_api.rs +++ b/golem-worker-executor/src/durable_host/golem/retry_api.rs @@ -128,8 +128,8 @@ impl Host for DurableWorkerCtx { let named_policy: NamedRetryPolicy = policy.into(); - if self.state.snapshotting_mode { - // Snapshot loading restores in-memory retry policy changes without recording them again. + if self.state.durability_is_suppressed() { + // Apply the in-memory change without creating durable history. } else if self.state.is_live() { self.public_state .worker() @@ -146,8 +146,8 @@ impl Host for DurableWorkerCtx { async fn remove_retry_policy(&mut self, name: String) -> anyhow::Result<()> { self.observe_function_call("golem::api::retry", "remove_retry_policy"); - if self.state.snapshotting_mode { - // Snapshot loading restores in-memory retry policy changes without recording them again. + if self.state.durability_is_suppressed() { + // Apply the in-memory change without creating durable history. } else if self.state.is_live() { self.public_state .worker() diff --git a/golem-worker-executor/src/durable_host/golem/v1x.rs b/golem-worker-executor/src/durable_host/golem/v1x.rs index 20c9ff95d6..c5229d17c5 100644 --- a/golem-worker-executor/src/durable_host/golem/v1x.rs +++ b/golem-worker-executor/src/durable_host/golem/v1x.rs @@ -322,7 +322,7 @@ impl Host for DurableWorkerCtx { async fn get_oplog_index(&mut self) -> anyhow::Result { self.observe_function_call("golem::api", "get_oplog_index"); - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { Ok(self.state.current_oplog_index().await.into()) } else if self.state.is_live() { // Use the index returned by `add` — a concurrently running host task (a durable @@ -362,7 +362,7 @@ impl Host for DurableWorkerCtx { oplog_idx: golem_api_1_x::oplog::OplogIndex, ) -> anyhow::Result<()> { self.observe_function_call("golem::api", "set_oplog_index"); - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { return Ok(()); } if self.state.is_live() { @@ -445,7 +445,7 @@ impl Host for DurableWorkerCtx { async fn mark_begin_operation(&mut self) -> anyhow::Result { self.observe_function_call("golem::api", "mark_begin_operation"); - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { Ok(self.state.current_oplog_index().await.into()) } else if self.state.is_live() { let next_idempotency_key_oplog_index = self @@ -532,7 +532,7 @@ impl Host for DurableWorkerCtx { begin: golem_api_1_x::oplog::OplogIndex, ) -> anyhow::Result<()> { self.observe_function_call("golem::api", "mark_end_operation"); - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { return Ok(()); } let begin_index = OplogIndex::from_u64(begin); @@ -1151,10 +1151,12 @@ impl Host for DurableWorkerCtx { .state .opens_durable_scope(&DurableFunctionType::WriteRemote) .then_some(oplog_index_cut_off); - self.public_state - .worker() - .commit_oplog_and_update_state(CommitLevel::Always) - .await; + if !self.is_unpersisted_execution() { + self.public_state + .worker() + .commit_oplog_and_update_state(CommitLevel::Always) + .await; + } let created_by = self.created_by(); let fork_result = loop { diff --git a/golem-worker-executor/src/durable_host/http/inline_retry.rs b/golem-worker-executor/src/durable_host/http/inline_retry.rs index 8ed5ffb12a..6477d34d55 100644 --- a/golem-worker-executor/src/durable_host/http/inline_retry.rs +++ b/golem-worker-executor/src/durable_host/http/inline_retry.rs @@ -127,7 +127,9 @@ pub(crate) fn take_http_background_retry_fallback( pub enum InlineRetryIneligible { /// Worker is in replay mode (not live). NotLive, - /// Worker is in snapshotting mode. + /// Worker is executing host operations without persistence. + UnpersistedExecution, + /// Worker is executing a snapshot load/save function. Snapshotting, /// Worker is inside a user-defined atomic region; a failure must escalate /// to trap+replay so the whole region re-executes. @@ -163,6 +165,9 @@ impl From for InlineRetryIneligible { fn from(reason: HttpRetryDisallowedReason) -> Self { match reason { HttpRetryDisallowedReason::NotLive => InlineRetryIneligible::NotLive, + HttpRetryDisallowedReason::UnpersistedExecution => { + InlineRetryIneligible::UnpersistedExecution + } HttpRetryDisallowedReason::Snapshotting => InlineRetryIneligible::Snapshotting, HttpRetryDisallowedReason::InAtomicRegion => InlineRetryIneligible::InAtomicRegion, HttpRetryDisallowedReason::NotIdempotent => InlineRetryIneligible::NotIdempotent, @@ -696,6 +701,7 @@ pub(crate) fn spawn_http_status_retry_after_body_finish, max_delay: Duration, begin_index: OplogIndex, + is_unpersisted_execution: bool, ) -> FutureIncomingResponseHandle { // No span: this task waits for the guest to finish its outgoing body, so its // duration is decided by guest code rather than by an operation the executor @@ -749,6 +755,7 @@ pub(crate) fn spawn_http_status_retry_after_body_finish( max_delay: Duration, begin_index: OplogIndex, execution_status: Arc>, + is_unpersisted_execution: bool, ) -> FutureIncomingResponseHandle { // Capture config fields individually since OutgoingRequestConfig is not Clone let use_tls = config.use_tls; @@ -998,6 +1006,7 @@ pub fn spawn_http_request_with_retry( max_in_function_retry_delay: max_delay, current_retry_policy_state, retry_properties, + is_unpersisted_execution, worker, }; @@ -1273,6 +1282,7 @@ pub async fn try_output_stream_inline_retry( exec_state.max_in_function_retry_delay, request_state.begin_index, ctx.execution_status.clone(), + ctx.is_unpersisted_execution(), ); HostFutureIncomingResponse::pending(retry_handle) } else { @@ -1550,7 +1560,7 @@ pub(crate) enum StatusRetryOutcome { /// This is invoked from `HostFutureIncomingResponse::get` *after* the response has /// arrived (so its status is known) but *before* the response is exposed to guest /// code or persisted. Behavior: -/// - In replay mode, snapshotting mode, or inside an atomic region +/// - In replay mode, unpersisted execution, snapshotting, or inside an atomic region /// the function is a no-op (`NoRetry`) — by design (atomic-region semantics: skip /// in v1, the user-land throw triggers atomic-region replay). /// - Otherwise eligibility is checked using the same rules as @@ -1821,6 +1831,7 @@ mod tests { DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: true, max_in_function_retry_delay: Duration::from_secs(1), } diff --git a/golem-worker-executor/src/durable_host/http/outgoing_http.rs b/golem-worker-executor/src/durable_host/http/outgoing_http.rs index fda84bcae9..6151effb29 100644 --- a/golem-worker-executor/src/durable_host/http/outgoing_http.rs +++ b/golem-worker-executor/src/durable_host/http/outgoing_http.rs @@ -100,6 +100,7 @@ pub(crate) async fn maybe_enable_http_background_retry( durable_state.max_in_function_retry_delay, state.begin_index, ctx.execution_status.clone(), + ctx.is_unpersisted_execution(), ); HostFutureIncomingResponse::pending(retry_handle) } else { @@ -184,6 +185,7 @@ pub(crate) async fn maybe_enable_http_pending_status_retry( agent_type, ctx.state.config.max_in_function_retry_delay, state.begin_index, + ctx.is_unpersisted_execution(), )) } else { old diff --git a/golem-worker-executor/src/durable_host/http/policy.rs b/golem-worker-executor/src/durable_host/http/policy.rs index 920860b08f..121033561f 100644 --- a/golem-worker-executor/src/durable_host/http/policy.rs +++ b/golem-worker-executor/src/durable_host/http/policy.rs @@ -65,7 +65,9 @@ pub(crate) fn is_http_request_idempotent( pub(crate) enum HttpRetryDisallowedReason { /// Worker is in replay mode (not live). NotLive, - /// Worker is in snapshotting mode. + /// Worker is executing host operations without persistence. + UnpersistedExecution, + /// Worker is executing a snapshot load/save function. Snapshotting, /// Worker is inside a user-defined atomic region; a failure must escalate /// to trap+replay so the whole region re-executes. @@ -85,6 +87,9 @@ pub(crate) fn http_worker_state_allows_retry( if !exec_state.is_live { return Err(HttpRetryDisallowedReason::NotLive); } + if exec_state.is_unpersisted_execution { + return Err(HttpRetryDisallowedReason::UnpersistedExecution); + } if exec_state.snapshotting_mode { return Err(HttpRetryDisallowedReason::Snapshotting); } @@ -297,6 +302,7 @@ mod tests { DurableExecutionState { is_live: true, snapshotting_mode: false, + is_unpersisted_execution: false, assume_idempotence: false, max_in_function_retry_delay: std::time::Duration::from_secs(1), } @@ -318,6 +324,16 @@ mod tests { ), Err(HttpRetryDisallowedReason::NotLive) ); + assert_eq!( + http_worker_state_allows_retry( + &DurableExecutionState { + is_unpersisted_execution: true, + ..live_exec_state() + }, + false + ), + Err(HttpRetryDisallowedReason::UnpersistedExecution) + ); assert_eq!( http_worker_state_allows_retry( &DurableExecutionState { diff --git a/golem-worker-executor/src/durable_host/http/types.rs b/golem-worker-executor/src/durable_host/http/types.rs index 0f11f82a37..d7309e4466 100644 --- a/golem-worker-executor/src/durable_host/http/types.rs +++ b/golem-worker-executor/src/durable_host/http/types.rs @@ -891,7 +891,7 @@ impl HostFutureIncomingResponse for DurableWorkerCtx { let handle = self_.rep(); let durable_execution_state = self.durable_execution_state(); - if durable_execution_state.is_live || self.state.snapshotting_mode { + if durable_execution_state.is_live { let request_state = self .state .open_http_requests @@ -1434,6 +1434,7 @@ impl DurableWorkerCtx { exec_state.max_in_function_retry_delay, request_state.begin_index, self.execution_status.clone(), + self.is_unpersisted_execution(), ); wasmtime_wasi_http::p2::types::HostFutureIncomingResponse::pending(retry_handle) } else { @@ -1529,7 +1530,7 @@ async fn persist_http_response( serializable_response: &SerializableHttpResponse, begin_index: golem_common::model::oplog::OplogIndex, ) { - if !ctx.state.snapshotting_mode { + if !ctx.state.durability_is_suppressed() { ctx.append_completed_child_call( HttpTypesFutureIncomingResponseGet::HOST_FUNCTION_NAME, &HostRequest::HttpRequest(request), diff --git a/golem-worker-executor/src/durable_host/logging/policy.rs b/golem-worker-executor/src/durable_host/logging/policy.rs index e39f9fa355..5b3983e69b 100644 --- a/golem-worker-executor/src/durable_host/logging/policy.rs +++ b/golem-worker-executor/src/durable_host/logging/policy.rs @@ -32,9 +32,10 @@ use std::sync::Arc; /// Applies the common log emission policy for a single worker log event. /// -/// `is_live` must be sampled from the worker state at the time of the call; `oplog` must be the -/// worker's private oplog (used by the [`LogEventEmitBehaviour::Always`] branch, which appends -/// without going through the invocation queue). +/// `is_live` and `is_unpersisted_execution` must be sampled from worker state at the time +/// of the call; `oplog` must be the worker's private oplog (used by the +/// [`LogEventEmitBehaviour::Always`] branch, which appends without going through the invocation +/// queue). pub async fn emit_log_event_with_state( event: InternalWorkerEvent, has_oplog_processor: bool, @@ -43,6 +44,7 @@ pub async fn emit_log_event_with_state( replay_state: &ReplayState, oplog: &Arc, is_live: bool, + is_unpersisted_execution: bool, ) { if let Some(entry) = event.as_oplog_entry() && let OplogEntry::Log { @@ -103,7 +105,9 @@ pub async fn emit_log_event_with_state( if !replay_state.seen_log(*level, context, message).await { // haven't seen this log before public_state.event_service().emit_event(event.clone(), true); - public_state.worker().add_to_oplog(entry).await; + if !is_unpersisted_execution { + public_state.worker().add_to_oplog(entry).await; + } } else { // we have persisted emitting this log before, so we mark it as non-live and // remove the entry from the seen log set. @@ -119,7 +123,10 @@ pub async fn emit_log_event_with_state( LogEventEmitBehaviour::Always => { public_state.event_service().emit_event(event.clone(), true); - if is_live && !replay_state.seen_log(*level, context, message).await { + if is_live + && !is_unpersisted_execution + && !replay_state.seen_log(*level, context, message).await + { oplog.add(entry).await; } } diff --git a/golem-worker-executor/src/durable_host/mod.rs b/golem-worker-executor/src/durable_host/mod.rs index 88084a4b8b..59ed5413d8 100644 --- a/golem-worker-executor/src/durable_host/mod.rs +++ b/golem-worker-executor/src/durable_host/mod.rs @@ -32,8 +32,13 @@ pub mod quota; mod random; pub mod rdbms; mod replay_state; +pub(crate) mod schema_value_stream; mod secrets; +pub use schema_value_stream::CoreTypesHost; mod sockets; +pub(crate) mod stream_bus; +pub(crate) mod stream_session; +pub(crate) mod stream_transport; mod suspendable_wait; pub mod tail_work; pub mod tool; @@ -46,14 +51,18 @@ use crate::durable_host::io::{ManagedStdErr, ManagedStdIn, ManagedStdOut}; use crate::durable_host::replay_state::{OplogEntryLookupResult, ReplayState}; use crate::metrics::ephemeral::record_non_suspending_failure; use crate::metrics::storage::{ - STORAGE_TYPE_FILESYSTEM, record_storage_bytes_deleted, record_storage_bytes_written, + STORAGE_TYPE_FILESYSTEM, record_filesystem_pool_released, record_storage_bytes_deleted, + record_storage_bytes_written, }; use crate::metrics::wasm::{record_number_of_replayed_functions, record_resume_worker}; use crate::model::event::InternalWorkerEvent; use crate::model::{ AgentConfig, ExecutionStatus, InvocationContext, LastError, ReadFileResult, TrapType, }; -use crate::services::active_workers::MemoryGrant; +use crate::services::active_workers::{ + FilesystemStoragePermit, MemoryGrant, bytes_to_filesystem_storage_permits, + filesystem_storage_permits_to_bytes, +}; use crate::services::agent_storage_meter::AgentStorageMeter; use crate::services::agent_types::AgentTypesService; use crate::services::agent_webhooks::AgentWebhooksService; @@ -176,6 +185,89 @@ impl WorkerDir { } } +#[derive(Debug)] +struct UnpersistedFilesystemStorage { + baseline_bytes: u64, + permit: Option, +} + +impl UnpersistedFilesystemStorage { + fn new(baseline_bytes: u64) -> Self { + Self { + baseline_bytes, + permit: None, + } + } + + fn extra_permits_for(&self, current_bytes: u64) -> u32 { + bytes_to_filesystem_storage_permits(current_bytes) + .saturating_sub(bytes_to_filesystem_storage_permits(self.baseline_bytes)) + } + + fn capacity_growth(&self, current_bytes: u64, new_bytes: u64) -> u64 { + let before = self.extra_permits_for(current_bytes); + let after = self.extra_permits_for(current_bytes.saturating_add(new_bytes)); + filesystem_storage_permits_to_bytes(after.saturating_sub(before)) + } + + fn merge(&mut self, permit: Option) { + let Some(permit) = permit else { + return; + }; + match &mut self.permit { + Some(existing) => existing.merge(permit), + None => self.permit = Some(permit), + } + } + + fn reconcile(&mut self, current_bytes: u64) { + let target = self.extra_permits_for(current_bytes) as usize; + let held = self + .permit + .as_ref() + .map_or(0, FilesystemStoragePermit::num_permits); + let excess = held.saturating_sub(target); + if excess > 0 + && let Some(permit) = &mut self.permit + { + record_filesystem_pool_released(filesystem_storage_permits_to_bytes(excess as u32)); + drop(permit.split(excess)); + } + if self + .permit + .as_ref() + .is_some_and(|permit| permit.num_permits() == 0) + { + self.permit = None; + } + } +} + +impl Drop for UnpersistedFilesystemStorage { + fn drop(&mut self) { + let permits = self + .permit + .as_ref() + .map_or(0, FilesystemStoragePermit::num_permits); + if permits > 0 { + record_filesystem_pool_released(filesystem_storage_permits_to_bytes(permits as u32)); + } + } +} + +/// An execution that performs real host operations but must not create replayable worker history. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum UnpersistedExecutionKind { + StreamingInvocation, +} + +#[derive(Debug, Clone, Copy)] +struct UnpersistedExecution { + kind: UnpersistedExecutionKind, +} + +type PreparedFilesystemStorageReservation = (Arc>, u64, bool); + impl Drop for WorkerDir { fn drop(&mut self) { if let WorkerDir::Deterministic(p) = self @@ -380,9 +472,18 @@ pub struct DurableWorkerCtx { resource_limits: Arc, linear_memory: LinearMemoryTracker, storage_meter: AgentStorageMeter, + /// Capacity acquired only for an unpersisted execution. Durable baseline capacity stays on + /// the running worker; this excess is Store-owned and is returned when the execution ends or + /// the Store is discarded. + unpersisted_filesystem_storage: Option, /// Per-instance cache of resolved typed guest export handles, populated /// lazily on first use during invocation dispatch. agent_export_funcs: AgentExportFuncs, + /// Source lifecycle for the live streaming invocation currently executing + /// in this Store. `schema-value-stream.wrap` attaches newly created source + /// endpoints to it so the invocation can keep the Store event loop alive + /// until downstream readers finish or detach. + live_stream_tracker: Option>, _store_alive_guard: StoreAliveGuard, } @@ -829,7 +930,9 @@ impl DurableWorkerCtx { resource_limits, linear_memory, storage_meter, + unpersisted_filesystem_storage: None, agent_export_funcs: AgentExportFuncs::default(), + live_stream_tracker: None, _store_alive_guard: StoreAliveGuard::new(), }) } @@ -1040,6 +1143,29 @@ impl DurableWorkerCtx { return Ok(()); } + if self.is_unpersisted_execution() { + for pending_event in self.pending_card_events_at_boundary().await? { + match pending_event.event { + QueuedCardEvent::Revoke(event) => { + self.apply_card_revoked(event.card_id, pending_event.oplog_index, true) + .await?; + } + QueuedCardEvent::Install(event) => { + let Some(card) = event.card else { + return Err(WorkerExecutorError::runtime( + "queued card install is missing card payload", + )); + }; + let _ = self + .apply_card_install(Some(pending_event.oplog_index), card) + .await?; + } + } + } + self.remove_expired_cards().await; + return Ok(()); + } + while let Some(pending_event) = self.pending_card_events_at_boundary().await?.first() { match &pending_event.event { QueuedCardEvent::Revoke(event) => { @@ -1155,7 +1281,9 @@ impl DurableWorkerCtx { _ => CardInstallFailure::NotFound, }; - if let Some(queued_event_index) = queued_event_index { + if let Some(queued_event_index) = queued_event_index + && !self.is_unpersisted_execution() + { self.public_state .worker() .add_and_commit_oplog(OplogEntry::card_install_failed( @@ -1180,10 +1308,12 @@ impl DurableWorkerCtx { .set_card_interest(self.owned_agent_id.clone(), &wallet_card_ids) .await; - self.public_state - .worker() - .add_and_commit_oplog(OplogEntry::card_installed(queued_event_index, card)) - .await; + if !self.is_unpersisted_execution() { + self.public_state + .worker() + .add_and_commit_oplog(OplogEntry::card_installed(queued_event_index, card)) + .await; + } Ok(Ok(())) } } @@ -1212,10 +1342,12 @@ impl DurableWorkerCtx { .set_card_interest(self.owned_agent_id.clone(), &wallet_card_ids) .await; - self.public_state - .worker() - .add_and_commit_oplog(OplogEntry::card_revoked(queued_event_index, card_id)) - .await; + if !self.is_unpersisted_execution() { + self.public_state + .worker() + .add_and_commit_oplog(OplogEntry::card_revoked(queued_event_index, card_id)) + .await; + } } Ok(()) @@ -1255,11 +1387,13 @@ impl DurableWorkerCtx { .set_card_interest(self.owned_agent_id.clone(), &wallet_card_ids) .await; - for card_id in cards_to_expire { - self.public_state - .worker() - .add_and_commit_oplog(OplogEntry::card_expired(card_id)) - .await; + if !self.is_unpersisted_execution() { + for card_id in cards_to_expire { + self.public_state + .worker() + .add_and_commit_oplog(OplogEntry::card_expired(card_id)) + .await; + } } } @@ -1436,6 +1570,27 @@ impl DurableWorkerCtx { if self.state.is_replay() { return Ok(()); } + if self.is_unpersisted_execution() { + let current_bytes = self.state.current_filesystem_storage_usage; + let capacity_bytes = self + .unpersisted_filesystem_storage + .as_ref() + .expect( + "unpersisted filesystem accounting must be active during unpersisted execution", + ) + .capacity_growth(current_bytes, new_bytes); + let permit = self + .public_state + .worker() + .acquire_unpersisted_filesystem_storage_space(capacity_bytes) + .await?; + self.unpersisted_filesystem_storage + .as_mut() + .expect("unpersisted filesystem accounting must remain active across acquisition") + .merge(permit); + self.state.current_filesystem_storage_usage = current_bytes.saturating_add(new_bytes); + return Ok(()); + } // Acquire the semaphore permit first (non-blocking try). Writing the // oplog entry after a confirmed acquire ensures the oplog accurately // reflects only committed storage changes — a failed acquire leaves no @@ -1474,6 +1629,16 @@ impl DurableWorkerCtx { if freed_bytes == 0 { return; } + if self.is_unpersisted_execution() { + self.state.current_filesystem_storage_usage -= freed_bytes; + self.unpersisted_filesystem_storage + .as_mut() + .expect( + "unpersisted filesystem accounting must be active during unpersisted execution", + ) + .reconcile(self.state.current_filesystem_storage_usage); + return; + } self.public_state .worker() .add_to_oplog(OplogEntry::filesystem_storage_usage_update( @@ -1510,13 +1675,28 @@ impl DurableWorkerCtx { pub(crate) fn prepare_filesystem_storage_reservation( &mut self, new_bytes: u64, - ) -> anyhow::Result>>> { + ) -> anyhow::Result>> { if new_bytes == 0 || self.state.is_replay() { return Ok(None); } self.check_filesystem_storage_quota(new_bytes)?; + let is_unpersisted_execution = self.is_unpersisted_execution(); + let capacity_bytes = if is_unpersisted_execution { + self.unpersisted_filesystem_storage + .as_ref() + .expect( + "unpersisted filesystem accounting must be active during unpersisted execution", + ) + .capacity_growth(self.state.current_filesystem_storage_usage, new_bytes) + } else { + new_bytes + }; self.state.current_filesystem_storage_usage += new_bytes; - Ok(Some(self.public_state.worker())) + Ok(Some(( + self.public_state.worker(), + capacity_bytes, + is_unpersisted_execution, + ))) } pub(crate) fn rollback_filesystem_storage_reservation(&mut self, new_bytes: u64) { @@ -1526,10 +1706,22 @@ impl DurableWorkerCtx { self.state.current_filesystem_storage_usage -= new_bytes; } - pub(crate) fn finish_filesystem_storage_reservation(&mut self, new_bytes: u64) { + pub(crate) fn finish_filesystem_storage_reservation( + &mut self, + new_bytes: u64, + is_unpersisted_execution: bool, + unpersisted_permit: Option, + ) { if new_bytes == 0 || self.state.is_replay() { return; } + if is_unpersisted_execution { + if let Some(storage) = &mut self.unpersisted_filesystem_storage { + storage.merge(unpersisted_permit); + } + return; + } + debug_assert!(unpersisted_permit.is_none()); let account_id = self.created_by().to_string(); let environment_id = self.state.owned_agent_id.environment_id().to_string(); record_storage_bytes_written( @@ -1552,14 +1744,27 @@ impl DurableWorkerCtx { None } else { self.state.current_filesystem_storage_usage -= freed_bytes; + if self.is_unpersisted_execution() { + self.unpersisted_filesystem_storage + .as_mut() + .expect("unpersisted filesystem accounting must be active during unpersisted execution") + .reconcile(self.state.current_filesystem_storage_usage); + } Some((self.public_state.worker(), freed_bytes)) } } - pub(crate) fn finish_filesystem_storage_release(&mut self, freed_bytes: u64) { + pub(crate) fn finish_filesystem_storage_release( + &mut self, + freed_bytes: u64, + is_unpersisted_execution: bool, + ) { if freed_bytes == 0 || self.state.is_replay() { return; } + if is_unpersisted_execution { + return; + } let account_id = self.created_by().to_string(); let environment_id = self.state.owned_agent_id.environment_id().to_string(); record_storage_bytes_deleted( @@ -1572,7 +1777,7 @@ impl DurableWorkerCtx { pub fn increase_memory(&mut self, delta: u64) { let (_, reconciling) = self.linear_memory.grow(delta, Instant::now()); - if self.state.is_live() && !reconciling { + if self.state.is_live() && !reconciling && !self.is_unpersisted_execution() { // This is called from the `memory.grow` async resource limiter, which // Wasmtime runs through a blocking libcall on the store's fiber. While // that libcall waits, the store cannot make progress, so nothing may be @@ -1616,7 +1821,11 @@ impl DurableWorkerCtx { crate::metrics::workers::record_worker_memory_grow_rejected(); return Err(GolemSpecificWasmTrap::WorkerOutOfMemory.into()); }; - tracker.retain_growth_grant(grant); + if self.is_unpersisted_execution() { + tracker.retain_transient_growth_grant(grant); + } else { + tracker.retain_growth_grant(grant); + } } Ok(true) } @@ -1916,6 +2125,7 @@ impl DurableWorkerCtx { &self.state.replay_state, &self.state.oplog, self.state.is_live(), + self.is_unpersisted_execution(), ) .await; } @@ -1924,7 +2134,7 @@ impl DurableWorkerCtx { &mut self, function_type: &DurableFunctionType, ) -> Result { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { let begin_index = self.state.current_oplog_index().await; self.state.current_retry_point = begin_index; return Ok(begin_index); @@ -2110,7 +2320,7 @@ impl DurableWorkerCtx { function_type: &DurableFunctionType, begin_index: OplogIndex, ) -> Result<(), WorkerExecutorError> { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { return Ok(()); } @@ -2285,7 +2495,7 @@ impl DurableWorkerCtx { where Err: From, { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { let (_, tx) = handler.create_new().await?; let begin_index = self.state.current_oplog_index().await; Ok((begin_index, tx)) @@ -2514,7 +2724,7 @@ impl DurableWorkerCtx { &mut self, begin_index: OplogIndex, ) -> Result<(), WorkerExecutorError> { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { Ok(()) } else if self.is_live() { // There is some logic in the test code that intercepts oplogs adds for _just_ the oplog the is provided to the worker. @@ -2543,7 +2753,7 @@ impl DurableWorkerCtx { &mut self, begin_index: OplogIndex, ) -> Result<(), WorkerExecutorError> { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { Ok(()) } else if self.is_live() { // There is some logic in the test code that intercepts oplogs adds for _just_ the oplog the is provided to the worker. @@ -2572,7 +2782,7 @@ impl DurableWorkerCtx { &mut self, begin_index: OplogIndex, ) -> Result<(), WorkerExecutorError> { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { return Ok(()); } else if self.is_live() { // There is some logic in the test code that intercepts oplogs adds for _just_ the oplog the is provided to the worker. @@ -2623,7 +2833,7 @@ impl DurableWorkerCtx { &mut self, begin_index: OplogIndex, ) -> Result<(), WorkerExecutorError> { - if self.state.snapshotting_mode { + if self.state.durability_is_suppressed() { return Ok(()); } else if self.is_live() { // There is some logic in the test code that intercepts oplogs adds for _just_ the oplog the is provided to the worker. @@ -3135,6 +3345,41 @@ impl DurableWorkerCtx { self.state.tail_work_tracker() } + pub(crate) fn live_stream_tracker(&self) -> Option> { + self.live_stream_tracker.clone() + } + + pub(crate) fn live_stream_event_capacity(&self) -> usize { + self.state + .config + .limits + .live_stream_event_broadcast_capacity + .get() + } + + pub(crate) fn is_live_streaming_invocation(&self) -> bool { + self.live_stream_tracker.is_some() + } + + pub(crate) fn is_unpersisted_execution(&self) -> bool { + self.state.is_unpersisted_execution() + } + + pub(crate) fn set_live_stream_tracker( + &mut self, + tracker: Arc, + ) { + assert!( + self.live_stream_tracker.is_none(), + "a live streaming invocation is already active in this Store" + ); + self.live_stream_tracker = Some(tracker); + } + + pub(crate) fn clear_live_stream_tracker(&mut self) { + self.live_stream_tracker = None; + } + /// Arms the optional per-invocation wall-clock deadline (`limits.max_invocation_duration`) /// and returns its guard. Called at the start of every guest invocation. /// @@ -3227,6 +3472,57 @@ impl DurableWorkerCtx { ) } + fn begin_unpersisted_execution(&mut self, kind: UnpersistedExecutionKind) { + assert!( + self.state.unpersisted_execution.is_none(), + "cannot begin {kind:?} while {:?} is active", + self.state + .unpersisted_execution + .map(|execution| execution.kind) + ); + assert!( + self.unpersisted_filesystem_storage.is_none(), + "unpersisted filesystem accounting is already active" + ); + + self.unpersisted_filesystem_storage = Some(UnpersistedFilesystemStorage::new( + self.state.current_filesystem_storage_usage, + )); + self.state.unpersisted_execution = Some(UnpersistedExecution { kind }); + } + + fn end_unpersisted_execution(&mut self, expected_kind: UnpersistedExecutionKind) { + let execution = self.state.unpersisted_execution.take().unwrap_or_else(|| { + panic!("cannot end {expected_kind:?}: no unpersisted execution is active") + }); + assert_eq!( + execution.kind, expected_kind, + "cannot end {expected_kind:?} while {:?} is active", + execution.kind + ); + + let storage = self + .unpersisted_filesystem_storage + .take() + .expect("unpersisted filesystem accounting must be active"); + self.state.current_filesystem_storage_usage = storage.baseline_bytes; + } + + pub(crate) fn begin_unpersisted_streaming_invocation(&mut self) { + self.begin_unpersisted_execution(UnpersistedExecutionKind::StreamingInvocation); + } + + pub(crate) fn end_unpersisted_streaming_invocation_if_active(&mut self) { + if matches!( + self.state + .unpersisted_execution + .map(|execution| execution.kind), + Some(UnpersistedExecutionKind::StreamingInvocation) + ) { + self.end_unpersisted_execution(UnpersistedExecutionKind::StreamingInvocation); + } + } + pub(crate) fn end_call_snapshotting_function_if_active(&mut self) { if self.state.snapshotting_mode { self.end_call_snapshotting_function(); @@ -3572,6 +3868,13 @@ impl StatusManagement for DurableWorkerCtx { return Some(*interrupt_kind); } } + if self + .live_stream_tracker + .as_ref() + .is_some_and(|tracker| tracker.cancellation_token().is_cancelled()) + { + return Some(InterruptKind::Interrupt(Timestamp::now_utc())); + } // An exceeded invocation or tail-work deadline surfaces as a synthetic interrupt so work // traps at the next epoch check. The corresponding invocation boundary converts the // unwind into the appropriate timeout failure. @@ -3654,7 +3957,7 @@ impl InvocationHooks for DurableWorkerCtx { &mut self, mut invocation: AgentInvocation, ) -> Result<(), WorkerExecutorError> { - if !self.state.snapshotting_mode { + if !self.state.durability_is_suppressed() { let stack = self.get_current_invocation_context().await; match &mut invocation { @@ -3696,11 +3999,12 @@ impl InvocationHooks for DurableWorkerCtx { full_function_name: &str, trap_type: &TrapType, ) -> RetryDecision { - self.cleanup_custom_durability_state(); + let is_unpersisted_execution = self.is_unpersisted_execution(); let current_idempotency_key = self.get_current_idempotency_key().await; if self.state.is_live() && !self.state.snapshotting_mode + && !is_unpersisted_execution && let Err(err) = concurrent::drain_queued_dropped_call_events(self).await { error!("failed to drain dropped durable calls before invocation failure entry: {err}"); @@ -3763,7 +4067,7 @@ impl InvocationHooks for DurableWorkerCtx { )), }; - if let Some(entry) = oplog_entry { + if !is_unpersisted_execution && let Some(entry) = oplog_entry { self.public_state.worker().add_and_commit_oplog(entry).await; }; @@ -3784,7 +4088,7 @@ impl InvocationHooks for DurableWorkerCtx { AgentStatus::Interrupted | AgentStatus::Exited ) || decision == RetryDecision::None; - if giving_up { + if giving_up && !is_unpersisted_execution { // Giving up, associating the stored result with the current and upcoming invocations if let Some(idempotency_key) = self.state.get_current_idempotency_key() { self.public_state @@ -3841,7 +4145,7 @@ impl InvocationHooks for DurableWorkerCtx { output: &mut AgentInvocationOutput, ) -> Result<(), WorkerExecutorError> { let is_live = self.state.is_live(); - if is_live && !self.state.snapshotting_mode { + if is_live && !self.state.snapshotting_mode && !self.is_unpersisted_execution() { concurrent::drain_queued_dropped_call_events(self) .await .map_err(|err| err.source)?; @@ -3856,7 +4160,7 @@ impl InvocationHooks for DurableWorkerCtx { } if is_live { - if !self.state.snapshotting_mode { + if !self.state.durability_is_suppressed() { let component_revision = output.component_revision.ok_or_else(|| { WorkerExecutorError::runtime( "component_revision missing in AgentInvocationOutput during replay", @@ -4009,7 +4313,7 @@ impl ResourceStore for DurableWorkerCtx { async fn add(&mut self, resource: ResourceAny, name: ResourceTypeId) -> u64 { let id = self.state.add(resource, name.clone()).await; let resource_id = AgentResourceId(id); - if self.state.is_live() { + if self.state.is_live() && !self.is_unpersisted_execution() { let entry = OplogEntry::create_resource(resource_id, name.clone()); self.public_state.worker().add_to_oplog(entry).await; } @@ -4020,7 +4324,7 @@ impl ResourceStore for DurableWorkerCtx { let result = self.state.borrow(resource_id).await; if let Some((resource_type_id, _)) = &result { let id = AgentResourceId(resource_id); - if self.state.is_live() { + if self.state.is_live() && !self.is_unpersisted_execution() { let entry = OplogEntry::drop_resource(id, resource_type_id.clone()); self.public_state.worker().add_to_oplog(entry).await; } @@ -4040,7 +4344,6 @@ impl UpdateManagement for DurableWorkerCtx { } fn begin_call_snapshotting_function(&mut self) { - // Snapshot load/save calls do not write durable host-call entries. if self.state.snapshotting_mode { warn!( "begin_call_snapshotting_function called while snapshotting is already active; \ @@ -4177,7 +4480,7 @@ impl InvocationContextManagement for DurableWorkerCtx { span.set_attribute(name.clone(), value.clone()); } - if is_live { + if is_live && !self.is_unpersisted_execution() { self.public_state .worker() .add_to_oplog(OplogEntry::StartSpan { @@ -4217,12 +4520,12 @@ impl InvocationContextManagement for DurableWorkerCtx { } async fn finish_span(&mut self, span_id: &SpanId) -> Result<(), WorkerExecutorError> { - if self.is_live() { + if self.is_live() && !self.is_unpersisted_execution() { self.public_state .worker() .add_to_oplog(OplogEntry::finish_span(span_id.clone())) .await; - } else { + } else if !self.is_live() { crate::get_oplog_entry!(self.state.replay_state, OplogEntry::FinishSpan)?; } @@ -4255,7 +4558,7 @@ impl InvocationContextManagement for DurableWorkerCtx { .invocation_context .set_attribute(span_id, key.to_string(), value.clone()) .map_err(WorkerExecutorError::runtime)?; - if self.is_live() { + if self.is_live() && !self.is_unpersisted_execution() { self.public_state .worker() .add_to_oplog(OplogEntry::set_span_attribute( @@ -4264,7 +4567,7 @@ impl InvocationContextManagement for DurableWorkerCtx { value, )) .await; - } else { + } else if !self.is_live() { crate::get_oplog_entry!(self.state.replay_state, OplogEntry::SetSpanAttribute)?; } Ok(()) @@ -4803,6 +5106,38 @@ mod tests { ); } + #[test] + fn unpersisted_filesystem_capacity_reuses_the_durable_baseline() { + let storage = UnpersistedFilesystemStorage::new(1500); + + assert_eq!(storage.capacity_growth(0, 1000), 0); + assert_eq!(storage.capacity_growth(0, 2200), 1024); + assert_eq!(storage.capacity_growth(1500, 600), 1024); + assert_eq!(storage.capacity_growth(2100, 300), 0); + } + + #[test] + async fn unpersisted_filesystem_capacity_is_released_before_store_reuse() { + let semaphore = crate::services::active_workers::FilesystemStorageSemaphore::new( + 4 * 1024, + Duration::from_millis(1), + ); + let mut storage = UnpersistedFilesystemStorage::new(500); + + let growth = storage.capacity_growth(500, 600); + storage.merge(semaphore.try_acquire(growth).await); + assert_eq!(semaphore.available_bytes(), 3 * 1024); + + storage.reconcile(500); + assert_eq!(semaphore.available_bytes(), 4 * 1024); + + let growth = storage.capacity_growth(500, 2000); + storage.merge(semaphore.try_acquire(growth).await); + assert_eq!(semaphore.available_bytes(), 2 * 1024); + drop(storage); + assert_eq!(semaphore.available_bytes(), 4 * 1024); + } + #[test] fn snapshot_boundary_all_clear_has_no_blocker() { assert_eq!(SnapshotBoundaryConditions::default().blocker(), None); @@ -4839,6 +5174,13 @@ mod tests { }, SnapshotBoundaryBlocker::Snapshotting, ), + ( + SnapshotBoundaryConditions { + streaming_invocation: true, + ..Default::default() + }, + SnapshotBoundaryBlocker::StreamingInvocation, + ), ( SnapshotBoundaryConditions { in_flight_host_call: true, @@ -4858,15 +5200,16 @@ mod tests { #[test] fn snapshot_boundary_any_condition_combination_blocks() { - // Exhaustive truth table over the five conditions: a snapshot is admitted iff every + // Exhaustive truth table over the six conditions: a snapshot is admitted iff every // condition is clear, and the reported blocker is always one of the set conditions. - for bits in 0u32..32 { + for bits in 0u32..64 { let conditions = SnapshotBoundaryConditions { replaying: bits & 1 != 0, open_atomic_region: bits & 2 != 0, open_durable_scope: bits & 4 != 0, snapshotting: bits & 8 != 0, - in_flight_host_call: bits & 16 != 0, + streaming_invocation: bits & 16 != 0, + in_flight_host_call: bits & 32 != 0, }; let blocker = conditions.blocker(); assert_eq!( @@ -4880,6 +5223,7 @@ mod tests { SnapshotBoundaryBlocker::OpenAtomicRegion => conditions.open_atomic_region, SnapshotBoundaryBlocker::OpenDurableScope => conditions.open_durable_scope, SnapshotBoundaryBlocker::Snapshotting => conditions.snapshotting, + SnapshotBoundaryBlocker::StreamingInvocation => conditions.streaming_invocation, SnapshotBoundaryBlocker::InFlightHostCall => conditions.in_flight_host_call, }; assert!( @@ -4898,18 +5242,18 @@ mod tests { let replaying = bits & 1 != 0; let open_atomic_region = bits & 2 != 0; let open_durable_scope = bits & 4 != 0; - let snapshotting = bits & 8 != 0; + let unpersisted_execution = bits & 8 != 0; assert_eq!( PrivateDurableWorkerState::clean_checkpoint_boundary( replaying, open_atomic_region, open_durable_scope, - snapshotting, + unpersisted_execution, ), bits == 0, "checkpoint must be admitted iff no condition is set; replaying: {replaying}, \ open_atomic_region: {open_atomic_region}, open_durable_scope: {open_durable_scope}, \ - snapshotting: {snapshotting}" + unpersisted_execution: {unpersisted_execution}" ); } } @@ -4918,19 +5262,20 @@ mod tests { fn snapshot_boundary_is_checkpoint_boundary_with_no_in_flight_host_call() { // The documented sync invariant between the two predicates: `blocker() == None` is // equivalent to `at_clean_checkpoint_boundary() && !has_in_flight_live_host_calls()`. - for bits in 0u32..32 { + for bits in 0u32..64 { let conditions = SnapshotBoundaryConditions { replaying: bits & 1 != 0, open_atomic_region: bits & 2 != 0, open_durable_scope: bits & 4 != 0, snapshotting: bits & 8 != 0, - in_flight_host_call: bits & 16 != 0, + streaming_invocation: bits & 16 != 0, + in_flight_host_call: bits & 32 != 0, }; let at_checkpoint_boundary = PrivateDurableWorkerState::clean_checkpoint_boundary( conditions.replaying, conditions.open_atomic_region, conditions.open_durable_scope, - conditions.snapshotting, + conditions.snapshotting || conditions.streaming_invocation, ); assert_eq!( conditions.blocker().is_none(), @@ -6383,6 +6728,7 @@ struct PrivateDurableWorkerState { /// the HttpRequestState. Keyed by outgoing request rep. pending_http_retry_eligibility: HashMap, + unpersisted_execution: Option, snapshotting_mode: bool, /// Tracks whether the currently executing invocation is restricted to read-only side effects. @@ -6727,6 +7073,7 @@ impl PrivateDurableWorkerState { open_filesystem_input_streams: HashSet::new(), file_stream_pollables: HashSet::new(), tcp_taken_streams: HashMap::new(), + unpersisted_execution: None, snapshotting_mode: false, invocation_strictness: InvocationStrictness::Normal, read_only_method_name: None, @@ -6862,13 +7209,13 @@ impl PrivateDurableWorkerState { /// [`DurableWorkerCtx::end_function`] — namely a non-idempotent remote write or the first /// (`None`) call of a batched remote write. /// - /// Snapshotting turns off persistence entirely, and `persist`/`replay` skip `end_function` - /// while snapshotting, so no scope must be opened either: otherwise the scope `Start` would be + /// Unpersisted execution turns off persistence entirely, and `persist`/`replay` skip + /// `end_function`, so no scope must be opened either: otherwise the scope `Start` would be /// committed with no matching `End`, corrupting later replay. - /// A snapshotting region never straddles a single scope's begin/end, so guarding both ends with - /// the same predicate keeps the durable-scope stack balanced. + /// An unpersisted execution never straddles a single scope's begin/end, so guarding both ends + /// with the same predicate keeps the durable-scope stack balanced. fn opens_durable_scope(&self, function_type: &DurableFunctionType) -> bool { - !self.snapshotting_mode + !self.durability_is_suppressed() && ((*function_type == DurableFunctionType::WriteRemote && !self.assume_idempotence) || matches!( *function_type, @@ -7215,6 +7562,14 @@ impl PrivateDurableWorkerState { self.replay_state.is_live() } + pub fn is_unpersisted_execution(&self) -> bool { + self.unpersisted_execution.is_some() + } + + fn durability_is_suppressed(&self) -> bool { + self.snapshotting_mode || self.is_unpersisted_execution() + } + /// Whether the current oplog tip is a structurally clean boundary at which a mid-invocation /// status checkpoint may be taken: we are live, no rollback-capable region is open (so no later /// trap/replay can append a jump that deletes the tip), and snapshotting is not active. The @@ -7225,7 +7580,7 @@ impl PrivateDurableWorkerState { !self.is_live(), !self.active_atomic_regions.is_empty(), !self.active_durable_scopes.is_empty(), - self.snapshotting_mode, + self.durability_is_suppressed(), ) } @@ -7236,9 +7591,9 @@ impl PrivateDurableWorkerState { replaying: bool, open_atomic_region: bool, open_durable_scope: bool, - snapshotting: bool, + unpersisted_execution: bool, ) -> bool { - !replaying && !open_atomic_region && !open_durable_scope && !snapshotting + !replaying && !open_atomic_region && !open_durable_scope && !unpersisted_execution } /// The first condition currently blocking a snapshot, or `None` when the worker is at a safe @@ -7262,6 +7617,10 @@ impl PrivateDurableWorkerState { open_atomic_region: !self.active_atomic_regions.is_empty(), open_durable_scope: !self.active_durable_scopes.is_empty(), snapshotting: self.snapshotting_mode, + streaming_invocation: matches!( + self.unpersisted_execution.map(|execution| execution.kind), + Some(UnpersistedExecutionKind::StreamingInvocation) + ), in_flight_host_call: self.has_in_flight_live_host_calls(), } .blocker() @@ -7318,6 +7677,8 @@ struct SnapshotBoundaryConditions { open_durable_scope: bool, /// A snapshotting function (save/load) call is already in progress. snapshotting: bool, + /// A live streaming invocation is currently using the Store without persistence. + streaming_invocation: bool, /// A live durable host call is in flight: its `Start` may precede the cut while its /// terminal entry lands after it. in_flight_host_call: bool, @@ -7335,6 +7696,8 @@ impl SnapshotBoundaryConditions { Some(SnapshotBoundaryBlocker::OpenDurableScope) } else if self.snapshotting { Some(SnapshotBoundaryBlocker::Snapshotting) + } else if self.streaming_invocation { + Some(SnapshotBoundaryBlocker::StreamingInvocation) } else if self.in_flight_host_call { Some(SnapshotBoundaryBlocker::InFlightHostCall) } else { @@ -7351,6 +7714,7 @@ pub enum SnapshotBoundaryBlocker { OpenAtomicRegion, OpenDurableScope, Snapshotting, + StreamingInvocation, InFlightHostCall, } @@ -7361,6 +7725,9 @@ impl Display for SnapshotBoundaryBlocker { Self::OpenAtomicRegion => write!(f, "an atomic region is still open"), Self::OpenDurableScope => write!(f, "a durable scope is still open"), Self::Snapshotting => write!(f, "a snapshot function call is already in progress"), + Self::StreamingInvocation => { + write!(f, "a live streaming invocation is already in progress") + } Self::InFlightHostCall => write!(f, "a durable host call is still in flight"), } } diff --git a/golem-worker-executor/src/durable_host/p3/cli.rs b/golem-worker-executor/src/durable_host/p3/cli.rs index ca167be9f8..a7b0a2000f 100644 --- a/golem-worker-executor/src/durable_host/p3/cli.rs +++ b/golem-worker-executor/src/durable_host/p3/cli.rs @@ -333,18 +333,26 @@ async fn emit_log_event_access( accessor: &Accessor>, event: InternalWorkerEvent, ) { - let (has_oplog_processor, owned_agent_id, public_state, replay_state, oplog, is_live) = - accessor.with(|mut access| { - let ctx = durable_worker_ctx::(access.data_mut()); - ( - ctx.state.component_metadata.metadata.has_oplog_processor(), - ctx.owned_agent_id.clone(), - ctx.public_state.clone(), - ctx.state.replay_state.clone(), - ctx.state.oplog.clone(), - ctx.state.is_live(), - ) - }); + let ( + has_oplog_processor, + owned_agent_id, + public_state, + replay_state, + oplog, + is_live, + is_unpersisted_execution, + ) = accessor.with(|mut access| { + let ctx = durable_worker_ctx::(access.data_mut()); + ( + ctx.state.component_metadata.metadata.has_oplog_processor(), + ctx.owned_agent_id.clone(), + ctx.public_state.clone(), + ctx.state.replay_state.clone(), + ctx.state.oplog.clone(), + ctx.state.is_live(), + ctx.is_unpersisted_execution(), + ) + }); logging_policy::emit_log_event_with_state::( event, @@ -354,6 +362,7 @@ async fn emit_log_event_access( &replay_state, &oplog, is_live, + is_unpersisted_execution, ) .await; } diff --git a/golem-worker-executor/src/durable_host/p3/filesystem.rs b/golem-worker-executor/src/durable_host/p3/filesystem.rs index 9b01694600..ab34ab0a58 100644 --- a/golem-worker-executor/src/durable_host/p3/filesystem.rs +++ b/golem-worker-executor/src/durable_host/p3/filesystem.rs @@ -572,26 +572,43 @@ where return Ok(()); } - if let Some(worker) = accessor + if let Some((worker, capacity_bytes, is_unpersisted_execution)) = accessor .with(|mut access| { durable_worker_ctx::(access.data_mut()) .prepare_filesystem_storage_reservation(bytes) }) .map_err(wasmtime::Error::from_anyhow)? { - if let Err(error) = worker.acquire_filesystem_storage_space(bytes).await { - accessor.with(|mut access| { - durable_worker_ctx::(access.data_mut()) - .rollback_filesystem_storage_reservation(bytes); - }); - return Err(wasmtime::Error::from_anyhow(error)); + let unpersisted_permit = match if is_unpersisted_execution { + worker + .acquire_unpersisted_filesystem_storage_space(capacity_bytes) + .await + } else { + worker + .acquire_filesystem_storage_space(capacity_bytes) + .await + .map(|()| None) + } { + Ok(permit) => permit, + Err(error) => { + accessor.with(|mut access| { + durable_worker_ctx::(access.data_mut()) + .rollback_filesystem_storage_reservation(bytes); + }); + return Err(wasmtime::Error::from_anyhow(error)); + } + }; + if !is_unpersisted_execution { + worker + .add_to_oplog(OplogEntry::filesystem_storage_usage_update(bytes as i64)) + .await; } - worker - .add_to_oplog(OplogEntry::filesystem_storage_usage_update(bytes as i64)) - .await; accessor.with(|mut access| { - durable_worker_ctx::(access.data_mut()) - .finish_filesystem_storage_reservation(bytes); + durable_worker_ctx::(access.data_mut()).finish_filesystem_storage_reservation( + bytes, + is_unpersisted_execution, + unpersisted_permit, + ); }); } @@ -610,16 +627,21 @@ where return Ok(()); } - if let Some((worker, bytes)) = accessor.with(|mut access| { - durable_worker_ctx::(access.data_mut()).prepare_filesystem_storage_release(bytes) + if let Some((worker, bytes, is_unpersisted_execution)) = accessor.with(|mut access| { + let ctx = durable_worker_ctx::(access.data_mut()); + let is_unpersisted_execution = ctx.is_unpersisted_execution(); + ctx.prepare_filesystem_storage_release(bytes) + .map(|(worker, bytes)| (worker, bytes, is_unpersisted_execution)) }) { - worker - .add_to_oplog(OplogEntry::filesystem_storage_usage_update(-(bytes as i64))) - .await; - worker.release_filesystem_storage_space(bytes).await; + if !is_unpersisted_execution { + worker + .add_to_oplog(OplogEntry::filesystem_storage_usage_update(-(bytes as i64))) + .await; + worker.release_filesystem_storage_space(bytes).await; + } accessor.with(|mut access| { durable_worker_ctx::(access.data_mut()) - .finish_filesystem_storage_release(bytes); + .finish_filesystem_storage_release(bytes, is_unpersisted_execution); }); } diff --git a/golem-worker-executor/src/durable_host/p3/http/replay.rs b/golem-worker-executor/src/durable_host/p3/http/replay.rs index f75415c7e0..ffa7e379e3 100644 --- a/golem-worker-executor/src/durable_host/p3/http/replay.rs +++ b/golem-worker-executor/src/durable_host/p3/http/replay.rs @@ -238,7 +238,10 @@ where Some(recorded_body) => { let (oplog, recording_enabled) = accessor.with(|mut access| { let ctx = durable_worker_ctx::(access.data_mut()); - (ctx.state.oplog.clone(), !ctx.state.snapshotting_mode) + ( + ctx.state.oplog.clone(), + !ctx.is_unpersisted_execution() && !ctx.state.snapshotting_mode, + ) }); if recording_enabled { drain_replayed_request_body_completing_recording( diff --git a/golem-worker-executor/src/durable_host/p3/http/send.rs b/golem-worker-executor/src/durable_host/p3/http/send.rs index 1d325e23ff..3e1a8b1d16 100644 --- a/golem-worker-executor/src/durable_host/p3/http/send.rs +++ b/golem-worker-executor/src/durable_host/p3/http/send.rs @@ -398,7 +398,7 @@ where let recording_enabled = store.with(|mut access| { !durable_worker_ctx::(access.data_mut()) .state - .snapshotting_mode + .durability_is_suppressed() }); let converted = match convert_physical_send_request::( store, @@ -507,7 +507,9 @@ where if (handle.trap_context().in_atomic_region || handle.is_observational()) && store.with(|mut access| { let ctx = durable_worker_ctx::(access.data_mut()); - ctx.state.is_live() && !ctx.state.snapshotting_mode + ctx.state.is_live() + && !ctx.is_unpersisted_execution() + && !ctx.state.snapshotting_mode }) && matching_status_retry_policy( store, @@ -1091,6 +1093,9 @@ where max_in_function_retry_delay, current_retry_policy_state, retry_properties, + is_unpersisted_execution: store.with(|mut access| { + durable_worker_ctx::(access.data_mut()).is_unpersisted_execution() + }), worker, } } diff --git a/golem-worker-executor/src/durable_host/quota/mod.rs b/golem-worker-executor/src/durable_host/quota/mod.rs index f5506f5646..b6449323d4 100644 --- a/golem-worker-executor/src/durable_host/quota/mod.rs +++ b/golem-worker-executor/src/durable_host/quota/mod.rs @@ -33,12 +33,14 @@ use golem_common::model::quota::ReserveResult; use golem_common::model::quota::ResourceName; use golem_common::model::{ScheduledAction, Timestamp}; use golem_schema::schema::schema_value::QuotaTokenValuePayload; -use golem_schema::schema::wit::wire::HostQuotaToken; +use golem_schema::schema::wit::wire::HostQuotaTokenWithStore; use golem_schema::schema::wit::{QuotaTokenHandleRep, QuotaTokenResolver}; use golem_service_base::error::worker_executor::GolemSpecificWasmTrap; use golem_service_base::error::worker_executor::WorkerExecutorError; use tracing::debug; -use wasmtime::component::Resource; +use wasmtime::component::{Accessor, Resource}; + +use crate::durable_host::schema_value_stream::CoreTypesHost; /// Borrow the [`QuotaTokenEntry`] stored inside a `quota-token` resource handle. /// @@ -446,11 +448,17 @@ impl Host for DurableWorkerCtx { /// opaque [`QuotaTokenHandleRep`] by golem-schema. The only operation the core /// interface declares for it is `drop`, which releases the underlying lease /// state back to the executor pool. -impl HostQuotaToken for DurableWorkerCtx { - async fn drop(&mut self, rep: Resource) -> anyhow::Result<()> { - DurabilityHost::observe_function_call(self, "golem::core::quota-token", "drop"); - self.table().delete(rep)?; - Ok(()) +impl HostQuotaTokenWithStore for CoreTypesHost { + async fn drop( + accessor: &Accessor, + rep: Resource, + ) -> anyhow::Result<()> { + accessor.with(|mut access| { + let ctx = access.get(); + DurabilityHost::observe_function_call(ctx, "golem::core::quota-token", "drop"); + ctx.table().delete(rep)?; + Ok(()) + }) } } diff --git a/golem-worker-executor/src/durable_host/schema_value_stream.rs b/golem-worker-executor/src/durable_host/schema_value_stream.rs new file mode 100644 index 0000000000..2e7b6b75ad --- /dev/null +++ b/golem-worker-executor/src/durable_host/schema_value_stream.rs @@ -0,0 +1,295 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::durable_host::DurableWorkerCtx; +use crate::durable_host::stream_transport::{ + LiveInputProducer, LiveStreamEndpoint, output_stream_pair, +}; +use crate::workerctx::WorkerCtx; +use golem_schema::schema::schema_value::{QuotaTokenValuePayload, SecretValuePayload}; +use golem_schema::schema::wit::wire::{ + Host, HostQuotaToken, HostSchemaValueStream, HostSchemaValueStreamWithStore, HostSecret, + HostWithStore, SchemaValueTree, Uuid, +}; +use golem_schema::schema::wit::{ + QuotaTokenHandleRep, QuotaTokenResolver, SchemaValueStreamResolver, SecretHandleRep, + SecretResolver, +}; +use golem_schema::schema::{SchemaValue, SchemaValueStream, SchemaValueStreamHandleRep}; +use golem_service_base::error::worker_executor::WorkerExecutorError; +use std::marker::PhantomData; +use wasmtime::StoreContextMut; +use wasmtime::component::{Accessor, HasData, Resource, StreamReader}; + +pub(crate) fn contains_stream(value: &SchemaValue) -> bool { + match value { + SchemaValue::Stream(_) => true, + SchemaValue::Record { fields } => fields.iter().any(contains_stream), + SchemaValue::Tuple { elements } + | SchemaValue::List { elements } + | SchemaValue::FixedList { elements } => elements.iter().any(contains_stream), + SchemaValue::Variant(payload) => payload.payload.as_deref().is_some_and(contains_stream), + SchemaValue::Map { entries } => entries + .iter() + .any(|(key, value)| contains_stream(key) || contains_stream(value)), + SchemaValue::Option { inner } => inner.as_deref().is_some_and(contains_stream), + SchemaValue::Result(payload) => match payload { + golem_schema::schema::schema_value::ResultValuePayload::Ok { value } + | golem_schema::schema::schema_value::ResultValuePayload::Err { value } => { + value.as_deref().is_some_and(contains_stream) + } + }, + SchemaValue::Union(payload) => contains_stream(&payload.body), + _ => false, + } +} + +pub struct StoreValueResolver<'a, 'store, Ctx: WorkerCtx> { + store: &'a mut StoreContextMut<'store, Ctx>, +} + +impl<'a, 'store, Ctx: WorkerCtx> StoreValueResolver<'a, 'store, Ctx> { + pub fn new(store: &'a mut StoreContextMut<'store, Ctx>) -> Self { + Self { store } + } +} + +impl QuotaTokenResolver for StoreValueResolver<'_, '_, Ctx> { + type Error = WorkerExecutorError; + + fn snapshot_handle( + &mut self, + handle: Resource, + ) -> Result { + self.store + .data_mut() + .durable_ctx_mut() + .snapshot_handle(handle) + } + + fn handle_from_snapshot( + &mut self, + snapshot: &QuotaTokenValuePayload, + ) -> Result, Self::Error> { + self.store + .data_mut() + .durable_ctx_mut() + .handle_from_snapshot(snapshot) + } + + fn drop_handle(&mut self, handle: Resource) { + self.store.data_mut().durable_ctx_mut().drop_handle(handle) + } +} + +impl SecretResolver for StoreValueResolver<'_, '_, Ctx> { + type Error = WorkerExecutorError; + + fn snapshot_secret_handle( + &mut self, + handle: Resource, + ) -> Result { + self.store + .data_mut() + .durable_ctx_mut() + .snapshot_secret_handle(handle) + } + + fn secret_handle_from_snapshot( + &mut self, + snapshot: &SecretValuePayload, + ) -> Result, Self::Error> { + self.store + .data_mut() + .durable_ctx_mut() + .secret_handle_from_snapshot(snapshot) + } + + fn drop_secret_handle(&mut self, handle: Resource) { + self.store + .data_mut() + .durable_ctx_mut() + .drop_secret_handle(handle) + } +} + +impl SchemaValueStreamResolver for StoreValueResolver<'_, '_, Ctx> { + type Error = WorkerExecutorError; + + fn handle_from_stream( + &mut self, + stream: SchemaValueStream, + ) -> Result, Self::Error> { + if let Some(tracker) = self.store.data().durable_ctx().live_stream_tracker() { + stream + .with_host_endpoint::(|endpoint| endpoint.attach(tracker)) + .map_err(WorkerExecutorError::runtime)?; + } + self.store + .data_mut() + .durable_ctx_mut() + .table() + .push(SchemaValueStreamHandleRep::new(stream)) + .map_err(|error| { + WorkerExecutorError::runtime(format!( + "failed to create schema-value-stream handle: {error}" + )) + }) + } + + fn stream_from_handle( + &mut self, + handle: Resource, + ) -> Result { + let stream = self + .store + .data_mut() + .durable_ctx_mut() + .table() + .delete(handle) + .map_err(|error| { + WorkerExecutorError::runtime(format!("invalid schema-value-stream handle: {error}")) + })? + .into_stream(); + Ok(stream) + } + + fn drop_stream_handle(&mut self, handle: Resource) { + let _ = self + .store + .data_mut() + .durable_ctx_mut() + .table() + .delete(handle); + } +} + +impl SchemaValueStreamResolver for DurableWorkerCtx { + type Error = WorkerExecutorError; + + fn handle_from_stream( + &mut self, + stream: SchemaValueStream, + ) -> Result, Self::Error> { + if let Some(tracker) = self.live_stream_tracker() { + stream + .with_host_endpoint::(|endpoint| endpoint.attach(tracker)) + .map_err(WorkerExecutorError::runtime)?; + } + self.table() + .push(SchemaValueStreamHandleRep::new(stream)) + .map_err(|error| { + WorkerExecutorError::runtime(format!( + "failed to create schema-value-stream handle: {error}" + )) + }) + } + + fn stream_from_handle( + &mut self, + handle: Resource, + ) -> Result { + self.table() + .delete(handle) + .map(SchemaValueStreamHandleRep::into_stream) + .map_err(|error| { + WorkerExecutorError::runtime(format!("invalid schema-value-stream handle: {error}")) + }) + } + + fn drop_stream_handle(&mut self, handle: Resource) { + let _ = self.table().delete(handle); + } +} + +pub struct CoreTypesHost(PhantomData); + +impl HasData for CoreTypesHost { + type Data<'a> = &'a mut DurableWorkerCtx; +} + +impl HostQuotaToken for DurableWorkerCtx {} +impl HostSecret for DurableWorkerCtx {} +impl HostSchemaValueStream for DurableWorkerCtx {} + +impl HostSchemaValueStreamWithStore for CoreTypesHost { + async fn wrap( + accessor: &Accessor, + reader: StreamReader, + ) -> anyhow::Result> { + accessor + .with(|mut access| -> wasmtime::Result<_> { + let tracker = access.get().live_stream_tracker(); + let capacity = access.get().live_stream_event_capacity(); + let (consumer, stream) = + output_stream_pair(tracker, capacity).map_err(wasmtime::Error::msg)?; + reader.pipe(&mut access, consumer)?; + access + .get() + .table() + .push(SchemaValueStreamHandleRep::new(stream)) + .map_err(|error| wasmtime::Error::msg(error.to_string())) + }) + .map_err(|error| anyhow::anyhow!(error.to_string())) + } + + async fn unwrap( + accessor: &Accessor, + value: Resource, + ) -> anyhow::Result> { + accessor + .with(|mut access| -> wasmtime::Result<_> { + let stream = access + .get() + .table() + .delete(value) + .map_err(|error| wasmtime::Error::msg(error.to_string())) + .map(SchemaValueStreamHandleRep::into_stream)?; + let endpoint = stream + .take_host_endpoint::() + .map_err(wasmtime::Error::msg)?; + StreamReader::new(&mut access, LiveInputProducer::new(endpoint)) + }) + .map_err(|error| anyhow::anyhow!(error.to_string())) + } + + async fn drop( + accessor: &Accessor, + rep: Resource, + ) -> anyhow::Result<()> { + accessor.with(|mut access| { + access + .get() + .table() + .delete(rep) + .map_err(|error| anyhow::anyhow!(error.to_string())) + })?; + Ok(()) + } +} + +impl Host for DurableWorkerCtx {} + +impl HostWithStore for CoreTypesHost { + async fn parse_uuid(_accessor: &Accessor, uuid: String) -> Result { + uuid::Uuid::parse_str(&uuid) + .map(Into::into) + .map_err(|error| error.to_string()) + } + + async fn uuid_to_string(_accessor: &Accessor, uuid: Uuid) -> String { + let uuid: uuid::Uuid = uuid.into(); + uuid.to_string() + } +} diff --git a/golem-worker-executor/src/durable_host/secrets/mod.rs b/golem-worker-executor/src/durable_host/secrets/mod.rs index 3e5da0e9a8..fc027baf8f 100644 --- a/golem-worker-executor/src/durable_host/secrets/mod.rs +++ b/golem-worker-executor/src/durable_host/secrets/mod.rs @@ -38,11 +38,13 @@ use golem_common::schema::schema_type::SchemaType; use golem_common::schema::schema_value::{SchemaValue, SecretValuePayload}; use golem_common::schema::validation::subtyping::is_equivalent_cross_graph; use golem_common::schema::validation::value::validate_value; -use golem_schema::schema::wit::wire::{HostSecret, SchemaValueTree}; +use golem_schema::schema::wit::wire::{HostSecretWithStore, SchemaValueTree}; use golem_schema::schema::wit::{SecretHandleRep, SecretResolver, decode_graph, encode_value_with}; use golem_service_base::error::worker_executor::WorkerExecutorError; use golem_service_base::model::agent_secret::AgentSecret; -use wasmtime::component::Resource; +use wasmtime::component::{Accessor, Resource}; + +use crate::durable_host::schema_value_stream::CoreTypesHost; fn secret_entry<'a, Ctx: WorkerCtx>( ctx: &'a mut DurableWorkerCtx, @@ -171,11 +173,17 @@ fn reveal_error_to_wit(error: SecretRevealError) -> SecretError { } } -impl HostSecret for DurableWorkerCtx { - async fn drop(&mut self, rep: Resource) -> anyhow::Result<()> { - DurabilityHost::observe_function_call(self, "golem::core::secret", "drop"); - self.table().delete(rep)?; - Ok(()) +impl HostSecretWithStore for CoreTypesHost { + async fn drop( + accessor: &Accessor, + rep: Resource, + ) -> anyhow::Result<()> { + accessor.with(|mut access| { + let ctx = access.get(); + DurabilityHost::observe_function_call(ctx, "golem::core::secret", "drop"); + ctx.table().delete(rep)?; + Ok(()) + }) } } diff --git a/golem-worker-executor/src/durable_host/stream_bus.rs b/golem-worker-executor/src/durable_host/stream_bus.rs new file mode 100644 index 0000000000..8f833b1634 --- /dev/null +++ b/golem-worker-executor/src/durable_host/stream_bus.rs @@ -0,0 +1,595 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use async_broadcast::{Receiver, RecvError, Sender, TrySendError, broadcast}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use tokio::sync::{Mutex, Notify}; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct LiveStreamEvent { + pub(crate) offset: u64, + pub(crate) payload: LiveStreamEventPayload, +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum LiveStreamEventPayload { + Item(T), + End, + Error(String), + Cancel(String), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LiveStreamBusCreateError { + ZeroCapacity, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LiveStreamPublishError { + Closed, + Terminated, + OffsetOverflow, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LiveStreamReceiveError { + Closed, + Lagged(u64), +} + +struct PublishState { + next_offset: u64, + terminated: bool, +} + +pub(crate) struct LiveStreamPublisher { + sender: Sender>, + state: Arc>, + primary_activated: Arc, + primary_changed: Arc, +} + +impl Clone for LiveStreamPublisher { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + state: self.state.clone(), + primary_activated: self.primary_activated.clone(), + primary_changed: self.primary_changed.clone(), + } + } +} + +impl LiveStreamPublisher { + pub(crate) async fn publish_item(&self, value: T) -> Result { + let mut state = self.state.lock().await; + if state.terminated { + return Err(LiveStreamPublishError::Terminated); + } + let following_offset = state + .next_offset + .checked_add(1) + .ok_or(LiveStreamPublishError::OffsetOverflow)?; + let offset = state.next_offset; + self.send_event(LiveStreamEvent { + offset, + payload: LiveStreamEventPayload::Item(value), + }) + .await?; + state.next_offset = following_offset; + Ok(offset) + } + + pub(crate) async fn publish_end(&self) -> Result { + self.publish_terminal(LiveStreamEventPayload::End).await + } + + pub(crate) async fn publish_error(&self, error: String) -> Result { + self.publish_terminal(LiveStreamEventPayload::Error(error)) + .await + } + + pub(crate) async fn publish_cancel( + &self, + reason: String, + ) -> Result { + self.publish_terminal(LiveStreamEventPayload::Cancel(reason)) + .await + } + + async fn publish_terminal( + &self, + payload: LiveStreamEventPayload, + ) -> Result { + let mut state = self.state.lock().await; + if state.terminated { + return Err(LiveStreamPublishError::Terminated); + } + let offset = state.next_offset; + self.send_event(LiveStreamEvent { offset, payload }).await?; + state.terminated = true; + Ok(offset) + } + + async fn send_event(&self, event: LiveStreamEvent) -> Result<(), LiveStreamPublishError> { + if self.primary_activated.load(Ordering::Acquire) { + self.sender + .broadcast(event) + .await + .map_err(|_| LiveStreamPublishError::Closed)?; + return Ok(()); + } + + match self.sender.try_broadcast(event) { + Ok(_) => Ok(()), + Err(TrySendError::Closed(_)) | Err(TrySendError::Inactive(_)) => { + Err(LiveStreamPublishError::Closed) + } + Err(TrySendError::Full(event)) => loop { + let changed = self.primary_changed.notified(); + if self.primary_activated.load(Ordering::Acquire) { + self.sender + .broadcast(event) + .await + .map_err(|_| LiveStreamPublishError::Closed)?; + return Ok(()); + } + if self.sender.is_closed() { + return Err(LiveStreamPublishError::Closed); + } + changed.await; + }, + } + } + + pub(crate) fn subscribe_tail(&self) -> AuxiliaryLiveStreamSubscriber { + AuxiliaryLiveStreamSubscriber { + receiver: self.sender.new_receiver(), + } + } + + pub(crate) fn close(&self) { + self.sender.close(); + } +} + +struct PrimaryDropGuard { + sender: Sender>, + on_drop: Arc, + primary_changed: Arc, + armed: bool, +} + +impl Drop for PrimaryDropGuard { + fn drop(&mut self) { + self.sender.close(); + self.primary_changed.notify_waiters(); + if self.armed { + (self.on_drop)(); + } + } +} + +pub(crate) struct ReservedPrimaryLiveStreamSubscriber { + receiver: Option>>, + drop_guard: Option>, + primary_activated: Arc, + primary_changed: Arc, +} + +impl ReservedPrimaryLiveStreamSubscriber { + pub(crate) fn activate(mut self) -> PrimaryLiveStreamSubscriber { + self.primary_activated.store(true, Ordering::Release); + self.primary_changed.notify_waiters(); + PrimaryLiveStreamSubscriber { + receiver: self + .receiver + .take() + .expect("reserved primary stream subscriber already activated"), + drop_guard: self + .drop_guard + .take() + .expect("reserved primary stream subscriber already activated"), + } + } +} + +pub(crate) struct PrimaryLiveStreamSubscriber { + receiver: Receiver>, + drop_guard: PrimaryDropGuard, +} + +impl PrimaryLiveStreamSubscriber { + pub(crate) async fn recv(&mut self) -> Result, LiveStreamReceiveError> { + let event = receive(&mut self.receiver).await?; + if !matches!(&event.payload, LiveStreamEventPayload::Item(_)) { + self.drop_guard.armed = false; + } + Ok(event) + } +} + +#[allow(dead_code)] +pub(crate) struct AuxiliaryLiveStreamSubscriber { + receiver: Receiver>, +} + +#[allow(dead_code)] +impl AuxiliaryLiveStreamSubscriber { + pub(crate) async fn recv(&mut self) -> Result, LiveStreamReceiveError> { + receive(&mut self.receiver).await + } +} + +async fn receive( + receiver: &mut Receiver>, +) -> Result, LiveStreamReceiveError> { + receiver.recv().await.map_err(|error| match error { + RecvError::Closed => LiveStreamReceiveError::Closed, + RecvError::Overflowed(missed) => LiveStreamReceiveError::Lagged(missed), + }) +} + +pub(crate) fn live_stream_bus( + capacity: usize, + on_primary_drop: impl Fn() + Send + Sync + 'static, +) -> Result< + ( + LiveStreamPublisher, + ReservedPrimaryLiveStreamSubscriber, + ), + LiveStreamBusCreateError, +> { + if capacity == 0 { + return Err(LiveStreamBusCreateError::ZeroCapacity); + } + let (mut sender, receiver) = broadcast(capacity); + sender.set_overflow(false); + let primary_activated = Arc::new(AtomicBool::new(false)); + let primary_changed = Arc::new(Notify::new()); + Ok(( + LiveStreamPublisher { + sender: sender.clone(), + state: Arc::new(Mutex::new(PublishState { + next_offset: 0, + terminated: false, + })), + primary_activated: primary_activated.clone(), + primary_changed: primary_changed.clone(), + }, + ReservedPrimaryLiveStreamSubscriber { + receiver: Some(receiver), + drop_guard: Some(PrimaryDropGuard { + sender, + on_drop: Arc::new(on_primary_drop), + primary_changed: primary_changed.clone(), + armed: true, + }), + primary_activated, + primary_changed, + }, + )) +} + +pub(crate) fn live_output_stream_bus( + capacity: usize, + invocation_cancellation: CancellationToken, +) -> Result< + ( + LiveStreamPublisher, + ReservedPrimaryLiveStreamSubscriber, + ), + LiveStreamBusCreateError, +> { + live_stream_bus(capacity, move || invocation_cancellation.cancel()) +} + +pub(crate) fn live_input_stream_bus( + capacity: usize, + stream_cancellation: CancellationToken, + producer_notification: Arc, +) -> Result< + ( + LiveStreamPublisher, + ReservedPrimaryLiveStreamSubscriber, + ), + LiveStreamBusCreateError, +> { + live_stream_bus(capacity, move || { + stream_cancellation.cancel(); + producer_notification.notify_one(); + }) +} + +#[cfg(test)] +mod tests { + use super::{ + LiveStreamBusCreateError, LiveStreamEvent, LiveStreamEventPayload, LiveStreamPublishError, + live_input_stream_bus, live_output_stream_bus, live_stream_bus, + }; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::Duration; + use test_r::test; + use tokio::sync::Notify; + use tokio_util::sync::CancellationToken; + + #[test] + fn rejects_zero_capacity() { + let result = live_stream_bus::(0, || {}); + + assert!(matches!( + result, + Err(LiveStreamBusCreateError::ZeroCapacity) + )); + } + + #[test] + async fn fans_out_ordered_events_with_identical_offsets() { + let (publisher, primary) = live_stream_bus(4, || {}).unwrap(); + let mut primary = primary.activate(); + let mut auxiliary = publisher.subscribe_tail(); + + assert_eq!(publisher.publish_item("first").await, Ok(0)); + assert_eq!(publisher.publish_item("second").await, Ok(1)); + assert_eq!(publisher.publish_end().await, Ok(2)); + + let expected = vec![ + LiveStreamEvent { + offset: 0, + payload: LiveStreamEventPayload::Item("first"), + }, + LiveStreamEvent { + offset: 1, + payload: LiveStreamEventPayload::Item("second"), + }, + LiveStreamEvent { + offset: 2, + payload: LiveStreamEventPayload::End, + }, + ]; + let mut primary_events = Vec::new(); + let mut auxiliary_events = Vec::new(); + for _ in 0..3 { + primary_events.push(primary.recv().await.unwrap()); + auxiliary_events.push(auxiliary.recv().await.unwrap()); + } + assert_eq!(primary_events, expected); + assert_eq!(auxiliary_events, expected); + } + + #[test] + async fn slowest_subscriber_applies_bounded_backpressure() { + let (publisher, primary) = live_stream_bus(1, || {}).unwrap(); + let mut primary = primary.activate(); + let mut auxiliary = publisher.subscribe_tail(); + + publisher.publish_item(1).await.unwrap(); + let blocked = tokio::spawn({ + let publisher = publisher.clone(); + async move { publisher.publish_item(2).await } + }); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(!blocked.is_finished()); + + assert_eq!(primary.recv().await.unwrap().offset, 0); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(!blocked.is_finished()); + assert_eq!(auxiliary.recv().await.unwrap().offset, 0); + assert_eq!(blocked.await.unwrap(), Ok(1)); + } + + #[test] + async fn late_subscriber_starts_at_the_current_tail() { + let (publisher, primary) = live_stream_bus(4, || {}).unwrap(); + let mut primary = primary.activate(); + publisher.publish_item(1).await.unwrap(); + let mut late = publisher.subscribe_tail(); + publisher.publish_item(2).await.unwrap(); + + assert_eq!(primary.recv().await.unwrap().offset, 0); + assert_eq!(primary.recv().await.unwrap().offset, 1); + assert_eq!(late.recv().await.unwrap().offset, 1); + } + + #[test] + async fn reserved_primary_buffers_one_item_and_backpressures_until_activation() { + let (publisher, primary) = live_stream_bus(1, || {}).unwrap(); + + assert_eq!(publisher.publish_item(1).await, Ok(0)); + let blocked = tokio::spawn({ + let publisher = publisher.clone(); + async move { publisher.publish_item(2).await } + }); + tokio::task::yield_now().await; + assert!(!blocked.is_finished()); + + let mut primary = primary.activate(); + assert_eq!(primary.recv().await.unwrap().offset, 0); + assert_eq!(blocked.await.unwrap(), Ok(1)); + assert_eq!(primary.recv().await.unwrap().offset, 1); + } + + #[test] + async fn reserved_primary_buffers_to_configured_capacity() { + let (publisher, primary) = live_stream_bus(3, || {}).unwrap(); + + for value in 1..=3 { + assert_eq!(publisher.publish_item(value).await, Ok(value - 1)); + } + let blocked = tokio::spawn({ + let publisher = publisher.clone(); + async move { publisher.publish_item(4).await } + }); + tokio::task::yield_now().await; + assert!(!blocked.is_finished()); + + let mut primary = primary.activate(); + for offset in 0..3 { + assert_eq!(primary.recv().await.unwrap().offset, offset); + } + assert_eq!(blocked.await.unwrap(), Ok(3)); + assert_eq!(primary.recv().await.unwrap().offset, 3); + } + + #[test] + async fn dropping_auxiliary_removes_only_its_backpressure() { + let (publisher, primary) = live_stream_bus(1, || {}).unwrap(); + let mut primary = primary.activate(); + let auxiliary = publisher.subscribe_tail(); + publisher.publish_item(1).await.unwrap(); + drop(auxiliary); + assert_eq!(primary.recv().await.unwrap().offset, 0); + + assert_eq!(publisher.publish_item(2).await, Ok(1)); + assert_eq!(primary.recv().await.unwrap().offset, 1); + } + + #[test] + async fn primary_loss_closes_the_bus_and_runs_its_scope_guard() { + let cancelled = Arc::new(AtomicBool::new(false)); + let (publisher, primary) = live_stream_bus(1, { + let cancelled = cancelled.clone(); + move || cancelled.store(true, Ordering::Release) + }) + .unwrap(); + drop(primary); + + assert!(cancelled.load(Ordering::Acquire)); + assert_eq!( + publisher.publish_item(1).await, + Err(LiveStreamPublishError::Closed) + ); + } + + #[test] + fn output_primary_loss_cancels_the_invocation() { + let invocation_cancellation = CancellationToken::new(); + let (_publisher, primary) = + live_output_stream_bus::(1, invocation_cancellation.clone()).unwrap(); + + drop(primary); + + assert!(invocation_cancellation.is_cancelled()); + } + + #[test] + async fn input_primary_loss_is_stream_scoped_and_notifies_the_producer() { + let invocation_cancellation = CancellationToken::new(); + let first_stream_cancellation = invocation_cancellation.child_token(); + let second_stream_cancellation = invocation_cancellation.child_token(); + let producer_notification = Arc::new(Notify::new()); + let (_first_publisher, first_primary) = live_input_stream_bus::( + 1, + first_stream_cancellation.clone(), + producer_notification.clone(), + ) + .unwrap(); + let (second_publisher, second_primary) = live_input_stream_bus( + 1, + second_stream_cancellation.clone(), + Arc::new(Notify::new()), + ) + .unwrap(); + let mut second_primary = second_primary.activate(); + let notification = producer_notification.notified(); + + drop(first_primary); + + notification.await; + assert!(first_stream_cancellation.is_cancelled()); + assert!(!second_stream_cancellation.is_cancelled()); + assert!(!invocation_cancellation.is_cancelled()); + assert_eq!(second_publisher.publish_item(1).await, Ok(0)); + assert_eq!(second_primary.recv().await.unwrap().offset, 0); + } + + #[test] + async fn stream_scoped_primary_loss_does_not_cancel_a_sibling_bus() { + let first_cancelled = Arc::new(AtomicBool::new(false)); + let second_cancelled = Arc::new(AtomicBool::new(false)); + let (_first_publisher, first_primary) = live_stream_bus::(1, { + let first_cancelled = first_cancelled.clone(); + move || first_cancelled.store(true, Ordering::Release) + }) + .unwrap(); + let (second_publisher, second_primary) = live_stream_bus(1, { + let second_cancelled = second_cancelled.clone(); + move || second_cancelled.store(true, Ordering::Release) + }) + .unwrap(); + let mut second_primary = second_primary.activate(); + + drop(first_primary); + assert!(first_cancelled.load(Ordering::Acquire)); + assert!(!second_cancelled.load(Ordering::Acquire)); + assert_eq!(second_publisher.publish_item(1).await, Ok(0)); + assert_eq!(second_primary.recv().await.unwrap().offset, 0); + } + + #[test] + async fn terminal_is_unique_and_rejects_every_later_publish() { + let cancelled = Arc::new(AtomicBool::new(false)); + let (publisher, primary) = live_stream_bus(4, { + let cancelled = cancelled.clone(); + move || cancelled.store(true, Ordering::Release) + }) + .unwrap(); + let mut primary = primary.activate(); + publisher.publish_item(1).await.unwrap(); + assert_eq!(publisher.publish_error("failed".to_string()).await, Ok(1)); + assert_eq!( + publisher.publish_end().await, + Err(LiveStreamPublishError::Terminated) + ); + assert_eq!( + publisher.publish_cancel("cancelled".to_string()).await, + Err(LiveStreamPublishError::Terminated) + ); + assert_eq!( + publisher.publish_item(2).await, + Err(LiveStreamPublishError::Terminated) + ); + + assert!(matches!( + primary.recv().await.unwrap().payload, + LiveStreamEventPayload::Item(1) + )); + assert!(matches!( + primary.recv().await.unwrap().payload, + LiveStreamEventPayload::Error(error) if error == "failed" + )); + drop(primary); + assert!(!cancelled.load(Ordering::Acquire)); + } + + #[test] + async fn checked_offset_overflow_does_not_publish() { + let (publisher, primary) = live_stream_bus(1, || {}).unwrap(); + let mut primary = primary.activate(); + publisher.state.lock().await.next_offset = u64::MAX; + + assert_eq!( + publisher.publish_item(1).await, + Err(LiveStreamPublishError::OffsetOverflow) + ); + assert!( + tokio::time::timeout(Duration::from_millis(20), primary.recv()) + .await + .is_err() + ); + } +} diff --git a/golem-worker-executor/src/durable_host/stream_session.rs b/golem-worker-executor/src/durable_host/stream_session.rs new file mode 100644 index 0000000000..a63edc7f05 --- /dev/null +++ b/golem-worker-executor/src/durable_host/stream_session.rs @@ -0,0 +1,2605 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::durable_host::schema_value_stream::contains_stream; +use crate::durable_host::stream_bus::{ + AuxiliaryLiveStreamSubscriber, LiveStreamEventPayload, LiveStreamPublishError, + LiveStreamPublisher, PrimaryLiveStreamSubscriber, +}; +#[cfg(test)] +use crate::durable_host::stream_transport::LiveStreamTracker; +use crate::durable_host::stream_transport::{ + LiveStreamEndpoint, LiveStreamPeer, SourceLifecycle, input_stream_pair, +}; +use golem_api_grpc::proto::golem::common::Empty; +use golem_api_grpc::proto::golem::schema::{ + FixedListValue, ListValue, MapEntry, MapValue, OptionValue, RecordValue, ResultValue, + SchemaValue as ProtoSchemaValue, SchemaValueStreamReference, TupleValue, UnionValue, + VariantValue, result_value as proto_result_value, schema_value as proto_schema_value, +}; +use golem_api_grpc::proto::golem::worker::{ + InputStreamAck, InputStreamEnd, InputStreamItem, InvocationRequest, InvocationResponse, + OutputStreamEnd, OutputStreamError, OutputStreamItem, StreamCancel, StreamCancelReason, + StreamCancelRole, input_stream_item, invocation_request, invocation_response, +}; +use golem_schema::schema::SchemaValue; +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use tokio::sync::mpsc; + +#[must_use = "imported stream registrations must be committed or rolled back"] +#[derive(Default)] +struct ImportedRegistrationBatch { + stream_ids: Vec, +} + +impl ImportedRegistrationBatch { + fn is_empty(&self) -> bool { + self.stream_ids.is_empty() + } +} + +#[derive(Clone)] +struct ImportedStreamRoute { + lifecycle: Arc, + publisher: LiveStreamPublisher, + next_sequence: Arc>, + acknowledgements: Arc>, + completed: tokio_util::sync::CancellationToken, +} + +#[derive(Default)] +struct InputAcknowledgementState { + next_offset: u64, + consumer_closed: bool, +} + +enum ImportedStreamState { + Registering(ImportedStreamRoute), + Active(ImportedStreamRoute), +} + +enum ImportedTerminal { + End, + Error(String), + Cancel(String), +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum InputItemAdmission { + Acknowledged(InputStreamAck), + ConsumerClosed, +} + +#[derive(Default)] +struct ImportedStreams { + states: HashMap, + cancelled: HashMap, +} + +struct ExportedStreamRoute { + lifecycle: Arc, + publisher: LiveStreamPublisher, + cancelled: tokio_util::sync::CancellationToken, + activated: tokio_util::sync::CancellationToken, + acknowledgements: Option>, + announced: Arc, + next_sent_offset: Arc, + terminal_sent: Arc, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SessionSide { + Client, + Server, +} + +#[derive(Clone)] +enum SessionFrames { + Requests(mpsc::Sender), + Responses(mpsc::Sender), +} + +enum OutboundStreamMessage { + Request(Box), + Response(Box), +} + +#[derive(Default)] +struct SessionActivity { + active: AtomicUsize, + changed: tokio::sync::Notify, +} + +impl SessionActivity { + fn start(self: &Arc) -> SessionActivityGuard { + self.active.fetch_add(1, Ordering::AcqRel); + self.changed.notify_waiters(); + SessionActivityGuard(self.clone()) + } + + fn is_idle(&self) -> bool { + self.active.load(Ordering::Acquire) == 0 + } + + async fn wait_until(&self, remaining: usize) { + loop { + let changed = self.changed.notified(); + if self.active.load(Ordering::Acquire) <= remaining { + return; + } + changed.await; + } + } +} + +struct SessionActivityGuard(Arc); + +impl Drop for SessionActivityGuard { + fn drop(&mut self) { + let previous = self.0.active.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "live invocation activity count underflow"); + self.0.changed.notify_waiters(); + } +} + +struct LiveValueSessionInner { + next_stream_id: std::sync::atomic::AtomicU64, + expected_remote_parity: u64, + side: SessionSide, + frames: SessionFrames, + exported: Mutex>, + imported: Mutex, + imported_changed: tokio::sync::Notify, + seen_remote_stream_ids: Mutex>, + activity: Arc, + cancelled: tokio_util::sync::CancellationToken, + failure: Mutex>, + stream_capacity: usize, +} + +/// Converts recursive live values to and from the session protocol. The two +/// peers allocate disjoint odd and even stream IDs, so sibling streams cannot +/// alias even when nested streams are discovered in later items. +#[derive(Clone)] +pub(crate) struct LiveValueSession { + inner: Arc, +} + +impl LiveValueSession { + #[cfg(test)] + pub(crate) fn new_client(frames: mpsc::Sender) -> Self { + Self::new_client_with_capacity(frames, 32) + } + + #[cfg(test)] + pub(crate) fn new_server(frames: mpsc::Sender) -> Self { + Self::new_server_with_capacity(frames, 32) + } + + pub(crate) fn new_client_with_capacity( + frames: mpsc::Sender, + stream_capacity: usize, + ) -> Self { + Self::new_with_capacity( + SessionSide::Client, + SessionFrames::Requests(frames), + stream_capacity, + ) + } + + pub(crate) fn new_server_with_capacity( + frames: mpsc::Sender, + stream_capacity: usize, + ) -> Self { + Self::new_with_capacity( + SessionSide::Server, + SessionFrames::Responses(frames), + stream_capacity, + ) + } + + fn new_with_capacity(side: SessionSide, frames: SessionFrames, stream_capacity: usize) -> Self { + assert!( + stream_capacity > 0, + "live stream bus capacity must be non-zero" + ); + let first_local_stream_id = match side { + SessionSide::Client => 1, + SessionSide::Server => 2, + }; + Self { + inner: Arc::new(LiveValueSessionInner { + next_stream_id: std::sync::atomic::AtomicU64::new(first_local_stream_id), + expected_remote_parity: 1 - (first_local_stream_id & 1), + side, + frames, + exported: Mutex::new(HashMap::new()), + imported: Mutex::new(ImportedStreams::default()), + imported_changed: tokio::sync::Notify::new(), + seen_remote_stream_ids: Mutex::new(HashSet::new()), + activity: Arc::new(SessionActivity::default()), + cancelled: tokio_util::sync::CancellationToken::new(), + failure: Mutex::new(None), + stream_capacity, + }), + } + } + + #[cfg(test)] + pub(crate) fn encode(&self, value: &SchemaValue) -> Result { + let (value, stream_ids) = self.encode_with_registered_streams(value)?; + self.activate_exported_streams(&stream_ids); + Ok(value) + } + + pub(crate) fn encode_pending( + &self, + value: &SchemaValue, + ) -> Result<(ProtoSchemaValue, Vec), String> { + self.encode_with_registered_streams(value) + } + + pub(crate) fn activate_exported_streams(&self, stream_ids: &[u64]) { + let exported = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned"); + for stream_id in stream_ids { + if let Some(route) = exported.get(stream_id) { + route.announced.store(true, Ordering::Release); + route.activated.cancel(); + } + } + } + + fn encode_with_registered_streams( + &self, + value: &SchemaValue, + ) -> Result<(ProtoSchemaValue, Vec), String> { + let mut registered_stream_ids = Vec::new(); + match self.encode_inner(value, &mut registered_stream_ids) { + Ok(value) => Ok((value, registered_stream_ids)), + Err(error) => { + self.discard_exported_streams(®istered_stream_ids); + Err(error) + } + } + } + + fn discard_exported_streams(&self, stream_ids: &[u64]) { + let mut exported = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned"); + for id in stream_ids { + if let Some(route) = exported.remove(id) { + route.cancelled.cancel(); + } + } + } + + fn encode_inner( + &self, + value: &SchemaValue, + registered_stream_ids: &mut Vec, + ) -> Result { + if !contains_stream(value) { + return value.clone().try_into(); + } + + let value = match value { + SchemaValue::Stream(stream) => { + let id = self.allocate_local_stream_id()?; + let endpoint = stream.take_host_endpoint::()?; + let lifecycle = endpoint.lifecycle(); + let publisher = endpoint.publisher(); + self.spawn_exported(id, endpoint.activate(), lifecycle, publisher)?; + registered_stream_ids.push(id); + proto_schema_value::Value::StreamReference(SchemaValueStreamReference { + stream_id: id, + }) + } + SchemaValue::Record { fields } => proto_schema_value::Value::RecordValue(RecordValue { + fields: fields + .iter() + .map(|field| self.encode_inner(field, registered_stream_ids)) + .collect::>()?, + }), + SchemaValue::Variant(value) => { + proto_schema_value::Value::VariantValue(Box::new(VariantValue { + case: value.case, + payload: value + .payload + .as_deref() + .map(|payload| self.encode_inner(payload, registered_stream_ids)) + .transpose()? + .map(Box::new), + })) + } + SchemaValue::Tuple { elements } => proto_schema_value::Value::TupleValue(TupleValue { + elements: elements + .iter() + .map(|element| self.encode_inner(element, registered_stream_ids)) + .collect::>()?, + }), + SchemaValue::List { elements } => proto_schema_value::Value::ListValue(ListValue { + elements: elements + .iter() + .map(|element| self.encode_inner(element, registered_stream_ids)) + .collect::>()?, + }), + SchemaValue::FixedList { elements } => { + proto_schema_value::Value::FixedListValue(FixedListValue { + elements: elements + .iter() + .map(|element| self.encode_inner(element, registered_stream_ids)) + .collect::>()?, + }) + } + SchemaValue::Map { entries } => proto_schema_value::Value::MapValue(MapValue { + entries: entries + .iter() + .map(|(key, value)| { + Ok(MapEntry { + key: Some(self.encode_inner(key, registered_stream_ids)?), + value: Some(self.encode_inner(value, registered_stream_ids)?), + }) + }) + .collect::>()?, + }), + SchemaValue::Option { inner } => { + proto_schema_value::Value::OptionValue(Box::new(OptionValue { + inner: inner + .as_deref() + .map(|inner| self.encode_inner(inner, registered_stream_ids)) + .transpose()? + .map(Box::new), + })) + } + SchemaValue::Result(result) => { + let result = match result { + golem_schema::schema::schema_value::ResultValuePayload::Ok { value } => { + match value.as_deref() { + Some(value) => proto_result_value::Result::Ok(Box::new( + self.encode_inner(value, registered_stream_ids)?, + )), + None => proto_result_value::Result::OkUnit(Empty {}), + } + } + golem_schema::schema::schema_value::ResultValuePayload::Err { value } => { + match value.as_deref() { + Some(value) => proto_result_value::Result::Err(Box::new( + self.encode_inner(value, registered_stream_ids)?, + )), + None => proto_result_value::Result::ErrUnit(Empty {}), + } + } + }; + proto_schema_value::Value::ResultValue(Box::new(ResultValue { + result: Some(result), + })) + } + SchemaValue::Union(value) => { + proto_schema_value::Value::UnionValue(Box::new(UnionValue { + tag: value.tag.clone(), + body: Some(Box::new( + self.encode_inner(&value.body, registered_stream_ids)?, + )), + })) + } + _ => { + return Err( + "a stream-bearing live value has an unsupported structural shape".to_string(), + ); + } + }; + Ok(ProtoSchemaValue { value: Some(value) }) + } + + pub(crate) async fn decode(&self, value: ProtoSchemaValue) -> Result { + self.decode_with_rollback(value, true).await + } + + pub(crate) async fn decode_start( + &self, + value: ProtoSchemaValue, + ) -> Result { + self.decode_with_rollback(value, false).await + } + + async fn decode_with_rollback( + &self, + value: ProtoSchemaValue, + notify_remote_on_rollback: bool, + ) -> Result { + let _activity = self.inner.activity.start(); + let (value, registrations) = self.decode_with_registered_streams(value); + match value { + Ok(value) => { + self.commit_imported_streams(registrations); + Ok(value) + } + Err(error) => { + if notify_remote_on_rollback { + self.rollback_imported_streams(registrations).await; + } else { + self.rollback_imported_streams_silently(registrations); + } + Err(error) + } + } + } + + fn decode_with_registered_streams( + &self, + value: ProtoSchemaValue, + ) -> (Result, ImportedRegistrationBatch) { + let mut registrations = ImportedRegistrationBatch::default(); + let value = self.decode_inner(value, &mut registrations); + (value, registrations) + } + + fn decode_inner( + &self, + value: ProtoSchemaValue, + registrations: &mut ImportedRegistrationBatch, + ) -> Result { + match value + .value + .ok_or_else(|| "schema value has no value".to_string())? + { + proto_schema_value::Value::StreamReference(reference) => { + self.validate_remote_id(reference.stream_id)?; + let (peer, stream) = + input_stream_pair(self.inner.stream_capacity, &self.inner.cancelled)?; + self.spawn_imported(reference.stream_id, peer)?; + registrations.stream_ids.push(reference.stream_id); + Ok(SchemaValue::Stream(stream)) + } + proto_schema_value::Value::RecordValue(value) => Ok(SchemaValue::Record { + fields: value + .fields + .into_iter() + .map(|field| self.decode_inner(field, registrations)) + .collect::>()?, + }), + proto_schema_value::Value::VariantValue(value) => Ok(SchemaValue::Variant( + golem_schema::schema::schema_value::VariantValuePayload { + case: value.case, + payload: value + .payload + .map(|payload| self.decode_inner(*payload, registrations).map(Box::new)) + .transpose()?, + }, + )), + proto_schema_value::Value::TupleValue(value) => Ok(SchemaValue::Tuple { + elements: value + .elements + .into_iter() + .map(|element| self.decode_inner(element, registrations)) + .collect::>()?, + }), + proto_schema_value::Value::ListValue(value) => Ok(SchemaValue::List { + elements: value + .elements + .into_iter() + .map(|element| self.decode_inner(element, registrations)) + .collect::>()?, + }), + proto_schema_value::Value::FixedListValue(value) => Ok(SchemaValue::FixedList { + elements: value + .elements + .into_iter() + .map(|element| self.decode_inner(element, registrations)) + .collect::>()?, + }), + proto_schema_value::Value::MapValue(value) => Ok(SchemaValue::Map { + entries: value + .entries + .into_iter() + .map(|entry| { + Ok(( + self.decode_inner( + entry + .key + .ok_or_else(|| "live map entry has no key".to_string())?, + registrations, + )?, + self.decode_inner( + entry + .value + .ok_or_else(|| "live map entry has no value".to_string())?, + registrations, + )?, + )) + }) + .collect::>()?, + }), + proto_schema_value::Value::OptionValue(value) => Ok(SchemaValue::Option { + inner: value + .inner + .map(|inner| self.decode_inner(*inner, registrations).map(Box::new)) + .transpose()?, + }), + proto_schema_value::Value::ResultValue(value) => { + let result = match value + .result + .ok_or_else(|| "result value has no result arm".to_string())? + { + proto_result_value::Result::Ok(value) => { + golem_schema::schema::schema_value::ResultValuePayload::Ok { + value: Some(Box::new(self.decode_inner(*value, registrations)?)), + } + } + proto_result_value::Result::Err(value) => { + golem_schema::schema::schema_value::ResultValuePayload::Err { + value: Some(Box::new(self.decode_inner(*value, registrations)?)), + } + } + proto_result_value::Result::OkUnit(_) => { + golem_schema::schema::schema_value::ResultValuePayload::Ok { value: None } + } + proto_result_value::Result::ErrUnit(_) => { + golem_schema::schema::schema_value::ResultValuePayload::Err { value: None } + } + }; + Ok(SchemaValue::Result(result)) + } + proto_schema_value::Value::UnionValue(value) => Ok(SchemaValue::Union( + golem_schema::schema::schema_value::UnionValuePayload { + tag: value.tag, + body: Box::new( + self.decode_inner( + *value + .body + .ok_or_else(|| "live union value has no body".to_string())?, + registrations, + )?, + ), + }, + )), + value => ProtoSchemaValue { value: Some(value) }.try_into(), + } + } + + pub(crate) async fn route_request( + &self, + request: invocation_request::Request, + ) -> Result { + let _activity = self.inner.activity.start(); + match request { + invocation_request::Request::InputItem(item) => { + if let InputItemAdmission::Acknowledged(ack) = self.admit_input_item(item).await? { + let route = self.imported_route(ack.stream_id)?; + let mut acknowledgements = route.acknowledgements.lock().await; + if acknowledgements.consumer_closed { + return Ok(true); + } + let next_offset = ack + .sequence + .checked_add(ack.logical_item_count) + .ok_or_else(|| { + format!("input stream {} ACK offset overflow", ack.stream_id) + })?; + if !self + .send_outbound(OutboundStreamMessage::Response(Box::new( + invocation_response::Response::InputAck(ack), + ))) + .await + { + return Err( + "invocation response stream closed before input ACK".to_string() + ); + } + acknowledgements.next_offset = next_offset; + } + Ok(true) + } + invocation_request::Request::InputEnd(end) => { + self.terminate_imported(end.stream_id, end.offset, ImportedTerminal::End) + .await?; + Ok(true) + } + invocation_request::Request::StreamCancel(cancel) => { + self.route_stream_cancel(cancel).await?; + Ok(true) + } + invocation_request::Request::Start(_) + | invocation_request::Request::ResumeAttach(_) => Ok(false), + } + } + + pub(crate) async fn route_response( + &self, + response: invocation_response::Response, + ) -> Result { + let _activity = self.inner.activity.start(); + match response { + invocation_response::Response::OutputItem(item) => { + let value = item + .value + .ok_or_else(|| format!("output stream {} item has no value", item.stream_id))?; + self.admit_imported_value(item.stream_id, item.offset, value) + .await?; + Ok(true) + } + invocation_response::Response::OutputEnd(end) => { + self.terminate_imported(end.stream_id, end.offset, ImportedTerminal::End) + .await?; + Ok(true) + } + invocation_response::Response::OutputError(error) => { + self.terminate_imported( + error.stream_id, + error.offset, + ImportedTerminal::Error(error.details), + ) + .await?; + Ok(true) + } + invocation_response::Response::InputAck(ack) => { + let sender = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned") + .get(&ack.stream_id) + .and_then(|route| route.acknowledgements.clone()) + .ok_or_else(|| { + format!("acknowledgement for unknown input stream {}", ack.stream_id) + })?; + sender + .send(ack) + .await + .map_err(|_| "input stream is no longer awaiting an ACK".to_string())?; + Ok(true) + } + invocation_response::Response::StreamCancel(cancel) => { + self.route_stream_cancel(cancel).await?; + Ok(true) + } + invocation_response::Response::Accepted(_) + | invocation_response::Response::Rejected(_) + | invocation_response::Response::Result(_) + | invocation_response::Response::AttachmentRevoked(_) + | invocation_response::Response::Finished(_) => Ok(false), + } + } + + async fn admit_imported_value( + &self, + stream_id: u64, + offset: u64, + value: ProtoSchemaValue, + ) -> Result<(), String> { + let route = self.imported_route(stream_id)?; + let mut next_sequence = route.next_sequence.lock().await; + if offset != *next_sequence { + return Err(format!( + "output stream {stream_id} expected offset {}, got {offset}", + *next_sequence + )); + } + let following_offset = offset + .checked_add(1) + .ok_or_else(|| format!("output stream {stream_id} offset overflow"))?; + let (value, registrations) = self.decode_with_registered_streams(value); + let value = match value { + Ok(value) => value, + Err(error) => { + self.rollback_imported_streams(registrations).await; + return Err(error); + } + }; + match route.publisher.publish_item(value).await { + Ok(published_offset) if published_offset == offset => { + self.commit_imported_streams(registrations); + *next_sequence = following_offset; + Ok(()) + } + Ok(published_offset) => { + self.rollback_imported_streams(registrations).await; + Err(format!( + "output stream {stream_id} published offset {published_offset}, expected {offset}" + )) + } + Err(error) => { + self.rollback_imported_streams(registrations).await; + Err(format!( + "failed to admit output stream {stream_id} item: {error:?}" + )) + } + } + } + + async fn terminate_imported( + &self, + stream_id: u64, + offset: u64, + terminal: ImportedTerminal, + ) -> Result<(), String> { + let route = self.imported_route(stream_id)?; + let next_sequence = *route.next_sequence.lock().await; + if offset != next_sequence { + return Err(format!( + "stream {stream_id} expected terminal offset {next_sequence}, got {offset}" + )); + } + if self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .cancelled + .contains_key(&stream_id) + { + self.remove_cancelled_imported(stream_id); + return Ok(()); + } + let result = match terminal { + ImportedTerminal::End => route.publisher.publish_end().await, + ImportedTerminal::Error(error) => route.publisher.publish_error(error).await, + ImportedTerminal::Cancel(details) => route.publisher.publish_cancel(details).await, + }; + if let Err(error) = result { + if error == LiveStreamPublishError::Closed + && self.remove_cancelled_imported(stream_id).is_some() + { + return Ok(()); + } + return Err(format!( + "failed to publish terminal for stream {stream_id}: {error:?}" + )); + } + route.lifecycle.finish(); + if self.remove_imported(stream_id).is_none() { + self.remove_cancelled_imported(stream_id); + } + Ok(()) + } + + async fn route_stream_cancel(&self, cancel: StreamCancel) -> Result<(), String> { + let role = StreamCancelRole::try_from(cancel.role) + .map_err(|_| format!("invalid stream cancellation role {}", cancel.role))?; + match (self.inner.side, role) { + (SessionSide::Server, StreamCancelRole::InputProducer) + | (SessionSide::Client, StreamCancelRole::OutputProducer) => { + self.terminate_imported( + cancel.stream_id, + cancel.offset, + ImportedTerminal::Cancel( + cancel + .details + .unwrap_or_else(|| "stream producer cancelled".to_string()), + ), + ) + .await + } + (SessionSide::Server, StreamCancelRole::OutputConsumer) => { + self.cancel_for_output_consumer(cancel).await + } + (SessionSide::Client, StreamCancelRole::InputConsumer) => { + self.cancel_exported(cancel.stream_id) + } + _ => Err(format!( + "unexpected {:?} stream cancellation for {:?} session", + role, self.inner.side + )), + } + } + + pub(crate) async fn admit_input_item( + &self, + item: InputStreamItem, + ) -> Result { + let stream_id = item.stream_id; + let route = self.imported_route(stream_id)?; + let mut next_sequence = route.next_sequence.lock().await; + if item.sequence != *next_sequence { + return Err(format!( + "input stream {stream_id} expected sequence {}, got {}", + *next_sequence, item.sequence + )); + } + + let logical_item_count = match item.payload.as_ref() { + Some(input_stream_item::Payload::Value(_)) => 1, + Some(input_stream_item::Payload::PackedU8(bytes)) if !bytes.is_empty() => { + u64::try_from(bytes.len()) + .map_err(|_| format!("input stream {stream_id} logical item count overflow"))? + } + Some(input_stream_item::Payload::PackedU8(_)) => { + return Err("packed-u8 input item must not be empty".to_string()); + } + None => return Err("input stream item has no payload".to_string()), + }; + let following_sequence = item + .sequence + .checked_add(logical_item_count) + .ok_or_else(|| format!("input stream {stream_id} sequence overflow"))?; + + let mut values = Vec::new(); + match item.payload { + Some(input_stream_item::Payload::Value(value)) => { + let (value, registrations) = self.decode_with_registered_streams(value); + match value { + Ok(value) => values.push((value, registrations)), + Err(error) => { + self.rollback_imported_streams(registrations).await; + return Err(error); + } + } + } + Some(input_stream_item::Payload::PackedU8(bytes)) if !bytes.is_empty() => { + values.extend( + bytes.into_iter().map(|value| { + (SchemaValue::U8(value), ImportedRegistrationBatch::default()) + }), + ); + } + Some(input_stream_item::Payload::PackedU8(_)) => { + return Err("packed-u8 input item must not be empty".to_string()); + } + None => return Err("input stream item has no payload".to_string()), + } + + let mut values = values.into_iter().enumerate(); + while let Some((index, (value, registrations))) = values.next() { + let expected_offset = item.sequence + index as u64; + match route.publisher.publish_item(value).await { + Ok(offset) if offset == expected_offset => { + self.commit_imported_streams(registrations); + } + Ok(offset) => { + self.rollback_imported_streams(registrations).await; + self.fail(format!( + "input stream {stream_id} published offset {offset}, expected {expected_offset}" + )); + return Err(format!( + "input stream {stream_id} published offset {offset}, expected {expected_offset}" + )); + } + Err(LiveStreamPublishError::Closed) => { + self.rollback_imported_streams(registrations).await; + for (_, (_, registrations)) in values { + self.rollback_imported_streams(registrations).await; + } + *next_sequence = following_sequence; + return Ok(InputItemAdmission::ConsumerClosed); + } + Err(error) => { + self.rollback_imported_streams(registrations).await; + return Err(format!( + "failed to admit input stream {stream_id} item: {error:?}" + )); + } + } + } + *next_sequence = following_sequence; + Ok(InputItemAdmission::Acknowledged(InputStreamAck { + stream_id, + sequence: item.sequence, + logical_item_count, + })) + } + + pub(crate) async fn wait_idle(&self) { + loop { + let activity_changed = self.inner.activity.changed.notified(); + let imported_changed = self.inner.imported_changed.notified(); + let imported_settled = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states + .values() + .all(|state| matches!(state, ImportedStreamState::Active(_))); + if self.inner.activity.is_idle() && imported_settled { + return; + } + tokio::select! { + _ = activity_changed => {} + _ = imported_changed => {} + } + } + } + + pub(crate) fn cancel(&self) { + self.inner.cancelled.cancel(); + let (imported, cancelled_imported) = { + let mut imported = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned"); + ( + std::mem::take(&mut imported.states), + std::mem::take(&mut imported.cancelled), + ) + }; + for state in imported.into_values() { + let route = match state { + ImportedStreamState::Registering(route) | ImportedStreamState::Active(route) => { + route + } + }; + route.publisher.close(); + route.lifecycle.finish(); + route.completed.cancel(); + } + for route in cancelled_imported.into_values() { + route.completed.cancel(); + } + let exported = std::mem::take( + &mut *self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned"), + ); + for route in exported.into_values() { + route.cancelled.cancel(); + } + self.inner.imported_changed.notify_waiters(); + } + + async fn cancel_for_output_consumer(&self, cancel: StreamCancel) -> Result<(), String> { + let imported = std::mem::take( + &mut self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states, + ); + let mut imported_terminals = Vec::with_capacity(imported.len()); + for (stream_id, state) in imported { + let route = match state { + ImportedStreamState::Registering(route) | ImportedStreamState::Active(route) => { + route + } + }; + imported_terminals.push((stream_id, route.next_sequence.clone())); + route.publisher.close(); + route.lifecycle.finish(); + route.completed.cancel(); + } + + let exported = std::mem::take( + &mut *self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned"), + ); + self.inner.cancelled.cancel(); + let mut exported_terminals = Vec::with_capacity(exported.len()); + for (stream_id, route) in exported { + exported_terminals.push(( + stream_id, + route.announced, + route.next_sent_offset, + route.terminal_sent, + )); + route.cancelled.cancel(); + } + self.inner.imported_changed.notify_waiters(); + + let frames = match &self.inner.frames { + SessionFrames::Responses(frames) => frames.clone(), + SessionFrames::Requests(_) => { + return Err( + "output-consumer cancellation is only valid for server sessions".into(), + ); + } + }; + self.inner.activity.wait_until(1).await; + imported_terminals.sort_unstable_by_key(|(stream_id, _)| *stream_id); + for (stream_id, next_sequence) in imported_terminals { + let offset = *next_sequence.lock().await; + if frames + .send(InvocationResponse { + response: Some(invocation_response::Response::StreamCancel(StreamCancel { + stream_id, + offset, + role: StreamCancelRole::InputConsumer as i32, + reason: cancel.reason, + details: cancel.details.clone(), + })), + }) + .await + .is_err() + { + return Ok(()); + } + } + + exported_terminals.sort_unstable_by_key(|(stream_id, _, _, _)| *stream_id); + for (stream_id, announced, next_sent_offset, terminal_sent) in exported_terminals { + if !announced.load(Ordering::Acquire) || terminal_sent.load(Ordering::Acquire) { + continue; + } + let offset = if stream_id == cancel.stream_id { + cancel.offset + } else { + next_sent_offset.load(Ordering::Acquire) + }; + if frames + .send(InvocationResponse { + response: Some(invocation_response::Response::StreamCancel(StreamCancel { + stream_id, + offset, + role: StreamCancelRole::OutputProducer as i32, + reason: cancel.reason, + details: cancel.details.clone(), + })), + }) + .await + .is_err() + { + return Ok(()); + } + } + Ok(()) + } + + pub(crate) fn is_cancelled(&self) -> bool { + self.inner.cancelled.is_cancelled() + } + + #[allow(dead_code)] + pub(crate) fn subscribe_output_tail( + &self, + stream_id: u64, + ) -> Result, String> { + self.inner + .exported + .lock() + .expect("live stream map mutex poisoned") + .get(&stream_id) + .map(|route| route.publisher.subscribe_tail()) + .ok_or_else(|| format!("subscription for unknown output stream {stream_id}")) + } + + pub(crate) async fn finish_invocation(&self) -> Result<(), String> { + loop { + let activity_changed = self.inner.activity.changed.notified(); + let imported_changed = self.inner.imported_changed.notified(); + let (imported, imported_settling) = { + let imported = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned"); + let mut active = Vec::new(); + let mut settling = false; + for (id, state) in &imported.states { + match state { + ImportedStreamState::Active(route) + if route.lifecycle.finished.load(Ordering::Acquire) => + { + settling = true; + } + ImportedStreamState::Registering(_) => settling = true, + ImportedStreamState::Active(_) => active.push(*id), + } + } + (active, settling) + }; + let (exported, exported_settling) = { + let exported = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned"); + let mut active = Vec::new(); + let mut settling = false; + for (id, route) in exported.iter() { + if route.lifecycle.finished.load(Ordering::Acquire) { + settling = true; + } else { + active.push(*id); + } + } + (active, settling) + }; + if imported.is_empty() && exported.is_empty() { + if !self.inner.activity.is_idle() || imported_settling || exported_settling { + tokio::select! { + _ = activity_changed => continue, + _ = imported_changed => continue, + } + } + return Ok(()); + } + + let details = format!( + "live invocation terminated with open imported streams {imported:?} and open exported streams {exported:?}" + ); + self.terminate_for_failure(&details).await; + return Err(details); + } + } + + pub(crate) async fn terminate_for_failure(&self, details: &str) { + let exported = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned") + .values() + .map(|route| route.publisher.clone()) + .collect::>(); + for publisher in exported { + let _ = publisher.publish_error(details.to_string()).await; + } + + let imported = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states + .keys() + .copied() + .collect::>(); + for id in imported { + self.cancel_imported(id, details).await; + } + + self.wait_idle().await; + *self + .inner + .failure + .lock() + .expect("live invocation failure mutex poisoned") = Some(details.to_string()); + self.cancel(); + } + + pub(crate) fn fail(&self, error: String) { + *self + .inner + .failure + .lock() + .expect("live invocation failure mutex poisoned") = Some(error); + self.cancel(); + } + + async fn send_outbound(&self, message: OutboundStreamMessage) -> bool { + match (&self.inner.frames, message) { + (SessionFrames::Requests(frames), OutboundStreamMessage::Request(request)) => { + tokio::select! { + result = frames.send(InvocationRequest { request: Some(*request) }) => result.is_ok(), + _ = self.inner.cancelled.cancelled() => false, + } + } + (SessionFrames::Responses(frames), OutboundStreamMessage::Response(response)) => { + tokio::select! { + result = frames.send(InvocationResponse { response: Some(*response) }) => result.is_ok(), + _ = self.inner.cancelled.cancelled() => false, + } + } + (SessionFrames::Requests(_), OutboundStreamMessage::Response(_)) + | (SessionFrames::Responses(_), OutboundStreamMessage::Request(_)) => { + self.fail( + "live session attempted to send a message in the wrong direction".to_string(), + ); + false + } + } + } + + async fn rollback_imported_streams(&self, registrations: ImportedRegistrationBatch) { + for id in registrations.stream_ids { + self.cancel_imported(id, "recursive stream registration was rolled back") + .await; + } + } + + fn rollback_imported_streams_silently(&self, registrations: ImportedRegistrationBatch) { + for id in registrations.stream_ids { + if let Some(route) = self.remove_imported(id) { + route.publisher.close(); + route.lifecycle.finish(); + } + } + } + + async fn cancel_imported(&self, id: u64, details: &str) { + let Some(route) = self.retain_cancelled_imported(id) else { + return; + }; + let mut acknowledgements = route.acknowledgements.lock().await; + acknowledgements.consumer_closed = true; + let offset = acknowledgements.next_offset; + let role = match self.inner.side { + SessionSide::Client => StreamCancelRole::OutputConsumer, + SessionSide::Server => StreamCancelRole::InputConsumer, + }; + let cancel = StreamCancel { + stream_id: id, + offset, + role: role as i32, + reason: StreamCancelReason::Cancelled as i32, + details: Some(details.to_string()), + }; + let message = match self.inner.side { + SessionSide::Client => OutboundStreamMessage::Request(Box::new( + invocation_request::Request::StreamCancel(cancel), + )), + SessionSide::Server => OutboundStreamMessage::Response(Box::new( + invocation_response::Response::StreamCancel(cancel), + )), + }; + let _ = self.send_outbound(message).await; + if self.inner.side == SessionSide::Client { + self.cancel(); + } + } + + fn remove_imported(&self, id: u64) -> Option { + let state = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states + .remove(&id); + self.inner.imported_changed.notify_waiters(); + state.map(|state| match state { + ImportedStreamState::Registering(route) | ImportedStreamState::Active(route) => { + route.completed.cancel(); + route + } + }) + } + + fn retain_cancelled_imported(&self, id: u64) -> Option { + let mut imported = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned"); + let state = imported.states.remove(&id)?; + let route = match state { + ImportedStreamState::Registering(route) | ImportedStreamState::Active(route) => route, + }; + route.completed.cancel(); + route.publisher.close(); + route.lifecycle.finish(); + imported.cancelled.insert(id, route.clone()); + drop(imported); + self.inner.imported_changed.notify_waiters(); + Some(route) + } + + fn remove_cancelled_imported(&self, id: u64) -> Option { + let route = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .cancelled + .remove(&id); + self.inner.imported_changed.notify_waiters(); + route + } + + fn cancel_exported(&self, id: u64) -> Result<(), String> { + let route = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned") + .remove(&id) + .ok_or_else(|| format!("cancellation for unknown local stream {id}"))?; + route.cancelled.cancel(); + Ok(()) + } + + fn commit_imported_streams(&self, registrations: ImportedRegistrationBatch) { + if registrations.is_empty() { + return; + } + let mut imported = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned"); + for id in registrations.stream_ids { + let route = match imported.states.get(&id) { + Some(ImportedStreamState::Registering(route)) => Some(route.clone()), + _ => None, + }; + if let Some(route) = route { + imported + .states + .insert(id, ImportedStreamState::Active(route)); + } + } + drop(imported); + self.inner.imported_changed.notify_waiters(); + } + + fn validate_remote_id(&self, id: u64) -> Result<(), String> { + if id == 0 || id & 1 != self.inner.expected_remote_parity { + Err(format!("invalid remote stream id {id}")) + } else { + Ok(()) + } + } + + fn imported_route(&self, id: u64) -> Result { + self.validate_remote_id(id)?; + let imported = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned"); + imported + .states + .get(&id) + .map(|state| match state { + ImportedStreamState::Registering(route) | ImportedStreamState::Active(route) => { + route.clone() + } + }) + .or_else(|| imported.cancelled.get(&id).cloned()) + .ok_or_else(|| format!("item for unknown remote stream {id}")) + } + + fn allocate_local_stream_id(&self) -> Result { + loop { + let current = self.inner.next_stream_id.load(Ordering::Acquire); + if current == 0 { + return Err("live stream ID space is exhausted".to_string()); + } + let next = current.checked_add(2).unwrap_or(0); + if self + .inner + .next_stream_id + .compare_exchange(current, next, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + return Ok(current); + } + } + } + + fn exported_end(&self, stream_id: u64, offset: u64) -> OutboundStreamMessage { + match self.inner.side { + SessionSide::Client => OutboundStreamMessage::Request(Box::new( + invocation_request::Request::InputEnd(InputStreamEnd { stream_id, offset }), + )), + SessionSide::Server => OutboundStreamMessage::Response(Box::new( + invocation_response::Response::OutputEnd(OutputStreamEnd { stream_id, offset }), + )), + } + } + + fn exported_error( + &self, + stream_id: u64, + offset: u64, + details: String, + ) -> OutboundStreamMessage { + match self.inner.side { + SessionSide::Client => OutboundStreamMessage::Request(Box::new( + invocation_request::Request::StreamCancel(StreamCancel { + stream_id, + offset, + role: StreamCancelRole::InputProducer as i32, + reason: StreamCancelReason::Cancelled as i32, + details: Some(details), + }), + )), + SessionSide::Server => OutboundStreamMessage::Response(Box::new( + invocation_response::Response::OutputError(OutputStreamError { + stream_id, + offset, + details, + }), + )), + } + } + + fn spawn_exported( + &self, + id: u64, + mut receiver: PrimaryLiveStreamSubscriber, + lifecycle: Arc, + publisher: LiveStreamPublisher, + ) -> Result<(), String> { + let cancelled = self.inner.cancelled.child_token(); + let activated = tokio_util::sync::CancellationToken::new(); + let announced = Arc::new(AtomicBool::new(false)); + let next_sent_offset = Arc::new(AtomicU64::new(0)); + let terminal_sent = Arc::new(AtomicBool::new(false)); + let (acknowledgements, mut acknowledgement_rx) = if self.inner.side == SessionSide::Client { + let (sender, receiver) = mpsc::channel(1); + (Some(sender), Some(receiver)) + } else { + (None, None) + }; + let mut exported = self + .inner + .exported + .lock() + .expect("live stream map mutex poisoned"); + match exported.entry(id) { + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(ExportedStreamRoute { + lifecycle, + publisher, + cancelled: cancelled.clone(), + activated: activated.clone(), + acknowledgements, + announced: announced.clone(), + next_sent_offset: next_sent_offset.clone(), + terminal_sent: terminal_sent.clone(), + }); + } + std::collections::hash_map::Entry::Occupied(_) => { + return Err(format!("duplicate local stream id {id}")); + } + } + drop(exported); + let session = self.clone(); + let activity = self.inner.activity.start(); + tokio::spawn(async move { + let _activity = activity; + tokio::select! { + _ = activated.cancelled() => {} + _ = cancelled.cancelled() => return, + _ = session.inner.cancelled.cancelled() => return, + } + loop { + let event = tokio::select! { + event = receiver.recv() => event, + _ = cancelled.cancelled() => break, + _ = session.inner.cancelled.cancelled() => break, + }; + let (message, registered_stream_ids, terminal, expected_ack, sent_offset) = + match event { + Ok(event) => { + match event.payload { + LiveStreamEventPayload::Item(value) => { + match session.encode_with_registered_streams(&value) { + Ok((value, registered_stream_ids)) => { + let message = match session.inner.side { + SessionSide::Client => OutboundStreamMessage::Request(Box::new( + invocation_request::Request::InputItem( + InputStreamItem { + stream_id: id, + sequence: event.offset, + payload: Some( + input_stream_item::Payload::Value(value), + ), + }, + ), + )), + SessionSide::Server => OutboundStreamMessage::Response(Box::new( + invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: id, + offset: event.offset, + value: Some(value), + }, + ), + )), + }; + ( + message, + registered_stream_ids, + false, + (session.inner.side == SessionSide::Client) + .then_some((event.offset, 1)), + Some(event.offset + 1), + ) + } + Err(error) => ( + session.exported_error(id, event.offset, error), + Vec::new(), + true, + None, + None, + ), + } + } + LiveStreamEventPayload::End => ( + session.exported_end(id, event.offset), + Vec::new(), + true, + None, + None, + ), + LiveStreamEventPayload::Error(error) + | LiveStreamEventPayload::Cancel(error) => ( + session.exported_error(id, event.offset, error), + Vec::new(), + true, + None, + None, + ), + } + } + Err(error) => ( + session.exported_error( + id, + 0, + format!("failed to receive live stream event: {error:?}"), + ), + Vec::new(), + true, + None, + None, + ), + }; + if !session.send_outbound(message).await { + session.discard_exported_streams(®istered_stream_ids); + break; + } + if let Some(sent_offset) = sent_offset { + next_sent_offset.store(sent_offset, Ordering::Release); + } + if terminal { + terminal_sent.store(true, Ordering::Release); + } + session.activate_exported_streams(®istered_stream_ids); + if let Some((sequence, logical_item_count)) = expected_ack { + let ack = tokio::select! { + ack = acknowledgement_rx + .as_mut() + .expect("client stream has no acknowledgement receiver") + .recv() => ack, + _ = cancelled.cancelled() => None, + _ = session.inner.cancelled.cancelled() => None, + }; + let Some(ack) = ack else { + session.discard_exported_streams(®istered_stream_ids); + break; + }; + if ack.stream_id != id + || ack.sequence != sequence + || ack.logical_item_count != logical_item_count + { + session.discard_exported_streams(®istered_stream_ids); + session.fail(format!( + "input stream {id} received invalid acknowledgement ({}, {})", + ack.sequence, ack.logical_item_count + )); + break; + } + } + if terminal { + break; + } + } + session + .inner + .exported + .lock() + .expect("live stream map mutex poisoned") + .remove(&id); + }); + Ok(()) + } + + fn spawn_imported(&self, id: u64, peer: LiveStreamPeer) -> Result<(), String> { + if !self + .inner + .seen_remote_stream_ids + .lock() + .expect("live stream ID set mutex poisoned") + .insert(id) + { + return Err(format!("duplicate remote stream id {id}")); + } + let completed = tokio_util::sync::CancellationToken::new(); + let previous = self + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states + .insert( + id, + ImportedStreamState::Registering(ImportedStreamRoute { + lifecycle: peer.lifecycle.clone(), + publisher: peer.publisher.clone(), + next_sequence: Arc::new(tokio::sync::Mutex::new(0)), + acknowledgements: Arc::new(tokio::sync::Mutex::new( + InputAcknowledgementState::default(), + )), + completed: completed.clone(), + }), + ); + debug_assert!(previous.is_none(), "new remote stream ID already imported"); + self.inner.imported_changed.notify_waiters(); + let session = self.clone(); + let activity = self.inner.activity.start(); + tokio::spawn(async move { + let _activity = activity; + tokio::select! { + _ = completed.cancelled() => {} + _ = session.inner.cancelled.cancelled() => { + peer.publisher.close(); + peer.lifecycle.finish(); + } + _ = peer.primary_dropped.notified() => { + session.cancel_imported(id, "stream consumer dropped its primary reader").await; + } + } + }); + Ok(()) + } +} + +#[cfg(test)] +mod bus_tests { + use super::*; + use crate::durable_host::stream_transport::{ + LiveStreamEndpoint, LiveStreamPeer, input_stream_pair, + }; + use golem_schema::schema::schema_value::{ + ResultValuePayload, UnionValuePayload, VariantValuePayload, + }; + use test_r::test; + use tokio_util::sync::CancellationToken; + + fn stream_id(value: &ProtoSchemaValue) -> u64 { + match value.value.as_ref() { + Some(proto_schema_value::Value::StreamReference(reference)) => reference.stream_id, + other => panic!("expected stream id, got {other:?}"), + } + } + + fn stream_reference(id: u64) -> ProtoSchemaValue { + ProtoSchemaValue { + value: Some(proto_schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: id }, + )), + } + } + + fn stream_source(capacity: usize) -> (LiveStreamPeer, SchemaValue) { + let cancellation = CancellationToken::new(); + let (peer, stream) = input_stream_pair(capacity, &cancellation).unwrap(); + (peer, SchemaValue::Stream(stream)) + } + + fn take_primary( + value: SchemaValue, + ) -> crate::durable_host::stream_bus::PrimaryLiveStreamSubscriber { + let SchemaValue::Stream(stream) = value else { + panic!("expected stream"); + }; + stream + .take_host_endpoint::() + .unwrap() + .activate() + } + + fn count_streams(value: &SchemaValue) -> usize { + match value { + SchemaValue::Stream(_) => 1, + SchemaValue::Record { fields } => fields.iter().map(count_streams).sum(), + SchemaValue::Variant(value) => value + .payload + .as_deref() + .map(count_streams) + .unwrap_or_default(), + SchemaValue::Tuple { elements } + | SchemaValue::List { elements } + | SchemaValue::FixedList { elements } => elements.iter().map(count_streams).sum(), + SchemaValue::Map { entries } => entries + .iter() + .map(|(key, value)| count_streams(key) + count_streams(value)) + .sum(), + SchemaValue::Option { inner } => { + inner.as_deref().map(count_streams).unwrap_or_default() + } + SchemaValue::Result(value) => match value { + ResultValuePayload::Ok { value } | ResultValuePayload::Err { value } => { + value.as_deref().map(count_streams).unwrap_or_default() + } + }, + SchemaValue::Union(value) => count_streams(&value.body), + _ => 0, + } + } + + #[test] + async fn recursive_composites_preserve_independent_streams() { + let mut sources = Vec::new(); + let mut next_stream = || { + let (source, stream) = stream_source(4); + sources.push(source); + stream + }; + let value = SchemaValue::Record { + fields: vec![ + SchemaValue::Variant(VariantValuePayload { + case: 1, + payload: Some(Box::new(next_stream())), + }), + SchemaValue::Tuple { + elements: vec![next_stream()], + }, + SchemaValue::List { + elements: vec![next_stream()], + }, + SchemaValue::FixedList { + elements: vec![next_stream()], + }, + SchemaValue::Map { + entries: vec![(next_stream(), next_stream())], + }, + SchemaValue::Option { + inner: Some(Box::new(next_stream())), + }, + SchemaValue::Result(ResultValuePayload::Ok { + value: Some(Box::new(next_stream())), + }), + SchemaValue::Result(ResultValuePayload::Err { + value: Some(Box::new(next_stream())), + }), + SchemaValue::Union(UnionValuePayload { + tag: "stream".to_string(), + body: Box::new(next_stream()), + }), + ], + }; + let (sender_frames, _sender_frame_rx) = mpsc::channel(32); + let sender = LiveValueSession::new_client(sender_frames); + let encoded = sender.encode(&value).unwrap(); + let (receiver_frames, _receiver_frame_rx) = mpsc::channel(32); + let receiver = LiveValueSession::new_server(receiver_frames); + + let decoded = receiver.decode(encoded).await.unwrap(); + + assert_eq!(count_streams(&decoded), 10); + assert_eq!(sources.len(), 10); + sender.cancel(); + receiver.cancel(); + } + + #[test] + async fn stream_ids_are_affine_checked_and_session_local() { + let (frames, _frame_rx) = mpsc::channel(8); + let receiver = LiveValueSession::new_client(frames); + assert_eq!( + receiver.decode(stream_reference(0)).await.unwrap_err(), + "invalid remote stream id 0" + ); + assert_eq!( + receiver.decode(stream_reference(1)).await.unwrap_err(), + "invalid remote stream id 1" + ); + assert!(matches!( + receiver.decode(stream_reference(2)).await.unwrap(), + SchemaValue::Stream(_) + )); + assert_eq!( + receiver.decode(stream_reference(2)).await.unwrap_err(), + "duplicate remote stream id 2" + ); + + let (_source, stream) = stream_source(4); + let SchemaValue::Stream(stream) = stream else { + unreachable!() + }; + let alias = stream.clone(); + let (frames, _frame_rx) = mpsc::channel(8); + let sender = LiveValueSession::new_client(frames); + sender.encode(&SchemaValue::Stream(stream)).unwrap(); + assert_eq!( + sender.encode(&SchemaValue::Stream(alias)).unwrap_err(), + "schema value stream was already transferred" + ); + receiver.cancel(); + sender.cancel(); + } + + #[test] + async fn recursive_registration_failures_roll_back_streams() { + let (_source, stream) = stream_source(4); + let SchemaValue::Stream(stream) = stream else { + unreachable!() + }; + let alias = stream.clone(); + let (frames, _frame_rx) = mpsc::channel(8); + let sender = LiveValueSession::new_client(frames); + assert!( + sender + .encode(&SchemaValue::Tuple { + elements: vec![SchemaValue::Stream(stream), SchemaValue::Stream(alias)], + }) + .is_err() + ); + sender.wait_idle().await; + + let (frames, mut frame_rx) = mpsc::channel(8); + let receiver = LiveValueSession::new_client(frames); + let error = receiver + .decode(ProtoSchemaValue { + value: Some(proto_schema_value::Value::TupleValue(TupleValue { + elements: vec![stream_reference(2), stream_reference(2)], + })), + }) + .await + .unwrap_err(); + assert_eq!(error, "duplicate remote stream id 2"); + let cancel = frame_rx.recv().await.unwrap().request; + assert!(matches!( + cancel, + Some(invocation_request::Request::StreamCancel(StreamCancel { + stream_id: 2, + role, + .. + })) if role == StreamCancelRole::OutputConsumer as i32 + )); + receiver.wait_idle().await; + } + + #[test] + async fn rejected_start_rolls_back_streams_without_pre_accept_events() { + let (frames, mut frame_rx) = mpsc::channel(8); + let receiver = LiveValueSession::new_server(frames); + + let error = receiver + .decode_start(ProtoSchemaValue { + value: Some(proto_schema_value::Value::TupleValue(TupleValue { + elements: vec![stream_reference(1), stream_reference(1)], + })), + }) + .await + .unwrap_err(); + + assert_eq!(error, "duplicate remote stream id 1"); + assert!(frame_rx.try_recv().is_err()); + assert!( + receiver + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states + .is_empty() + ); + receiver.cancel(); + } + + #[test] + fn local_stream_ids_are_not_reused_at_exhaustion() { + let (frames, _frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_client(frames); + let first_id = stream_id(&session.encode(&stream_source(1).1).unwrap()); + session + .inner + .next_stream_id + .store(u64::MAX, Ordering::Release); + let last_id = stream_id(&session.encode(&stream_source(1).1).unwrap()); + + assert_ne!(first_id, last_id); + assert_eq!( + session.encode(&stream_source(1).1).unwrap_err(), + "live stream ID space is exhausted" + ); + session.cancel(); + } + + #[test] + async fn first_output_can_arrive_before_the_session_registers_its_reader() { + let (source, stream) = stream_source(1); + assert_eq!( + source.publisher.publish_item(SchemaValue::U32(7)).await, + Ok(0) + ); + let (frames, mut frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_server_with_capacity(frames, 1); + let id = stream_id(&session.encode(&stream).unwrap()); + + let Some(invocation_response::Response::OutputItem(item)) = + frame_rx.recv().await.unwrap().response + else { + panic!("expected output item"); + }; + assert_eq!(item.stream_id, id); + assert_eq!(item.offset, 0); + assert_eq!( + SchemaValue::try_from(item.value.unwrap()).unwrap(), + SchemaValue::U32(7) + ); + session.cancel(); + } + + #[test] + async fn exported_stream_applies_bus_backpressure_and_preserves_order() { + let (source, stream) = stream_source(1); + let (frames, mut frame_rx) = mpsc::channel(1); + let session = LiveValueSession::new_server_with_capacity(frames, 1); + let id = stream_id(&session.encode(&stream).unwrap()); + source + .publisher + .publish_item(SchemaValue::String("first".to_string())) + .await + .unwrap(); + source + .publisher + .publish_item(SchemaValue::String("second".to_string())) + .await + .unwrap(); + source + .publisher + .publish_item(SchemaValue::String("third".to_string())) + .await + .unwrap(); + let fourth = tokio::spawn({ + let publisher = source.publisher.clone(); + async move { + publisher + .publish_item(SchemaValue::String("fourth".to_string())) + .await + } + }); + tokio::task::yield_now().await; + assert!(!fourth.is_finished()); + + for (offset, expected) in ["first", "second", "third", "fourth"] + .into_iter() + .enumerate() + { + let Some(invocation_response::Response::OutputItem(item)) = + frame_rx.recv().await.unwrap().response + else { + panic!("expected output item"); + }; + assert_eq!(item.stream_id, id); + assert_eq!(item.offset, offset as u64); + assert_eq!( + SchemaValue::try_from(item.value.unwrap()).unwrap(), + SchemaValue::String(expected.to_string()) + ); + } + assert_eq!(fourth.await.unwrap(), Ok(3)); + session.cancel(); + } + + #[test] + async fn output_consumer_cancellation_confirms_all_open_output_terminals() { + let (first_source, first_stream) = stream_source(4); + let (second_source, second_stream) = stream_source(4); + let (frames, mut frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_server(frames); + let encoded = session + .encode(&SchemaValue::Tuple { + elements: vec![first_stream, second_stream], + }) + .unwrap(); + let proto_schema_value::Value::TupleValue(tuple) = encoded.value.unwrap() else { + panic!("expected output stream tuple"); + }; + let first_id = stream_id(&tuple.elements[0]); + let second_id = stream_id(&tuple.elements[1]); + first_source + .publisher + .publish_item(SchemaValue::U32(1)) + .await + .unwrap(); + second_source + .publisher + .publish_item(SchemaValue::U32(2)) + .await + .unwrap(); + let mut item_streams = HashSet::new(); + while item_streams.len() < 2 { + let frame = frame_rx.recv().await.unwrap(); + if let Some(invocation_response::Response::OutputItem(item)) = frame.response { + item_streams.insert(item.stream_id); + } + } + + session + .route_request(invocation_request::Request::StreamCancel(StreamCancel { + stream_id: first_id, + offset: 0, + role: StreamCancelRole::OutputConsumer as i32, + reason: StreamCancelReason::Cancelled as i32, + details: Some("consumer stopped".to_string()), + })) + .await + .unwrap(); + + let mut terminals = HashMap::new(); + while terminals.len() < 2 { + let frame = frame_rx.recv().await.unwrap(); + if let Some(invocation_response::Response::StreamCancel(cancel)) = frame.response { + assert_eq!(cancel.role(), StreamCancelRole::OutputProducer); + terminals.insert(cancel.stream_id, cancel.offset); + } + } + assert_eq!(terminals.get(&first_id), Some(&0)); + assert_eq!(terminals.get(&second_id), Some(&1)); + session.wait_idle().await; + } + + #[test] + async fn auxiliary_output_subscription_starts_at_the_current_tail() { + let (source, stream) = stream_source(4); + let (frames, _frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_server_with_capacity(frames, 4); + let id = stream_id(&session.encode(&stream).unwrap()); + source + .publisher + .publish_item(SchemaValue::String("before".to_string())) + .await + .unwrap(); + let mut auxiliary = session.subscribe_output_tail(id).unwrap(); + source + .publisher + .publish_item(SchemaValue::String("after".to_string())) + .await + .unwrap(); + + let event = auxiliary.recv().await.unwrap(); + + assert_eq!(event.offset, 1); + assert_eq!( + event.payload, + LiveStreamEventPayload::Item(SchemaValue::String("after".to_string())) + ); + session.cancel(); + } + + #[test] + async fn imported_stream_prefetches_into_the_bus_and_terminates_once() { + let (frames, _frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_client_with_capacity(frames, 2); + let value = session.decode(stream_reference(2)).await.unwrap(); + let mut primary = take_primary(value); + + session + .route_response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 2, + offset: 0, + value: Some(SchemaValue::U32(42).try_into().unwrap()), + }, + )) + .await + .unwrap(); + let item = primary.recv().await.unwrap(); + assert_eq!(item.offset, 0); + assert_eq!( + item.payload, + LiveStreamEventPayload::Item(SchemaValue::U32(42)) + ); + + session + .route_response(invocation_response::Response::OutputEnd(OutputStreamEnd { + stream_id: 2, + offset: 1, + })) + .await + .unwrap(); + let end = primary.recv().await.unwrap(); + assert_eq!(end.offset, 1); + assert_eq!(end.payload, LiveStreamEventPayload::End); + assert_eq!( + session + .route_response(invocation_response::Response::OutputEnd(OutputStreamEnd { + stream_id: 2, + offset: 1, + })) + .await + .unwrap_err(), + "item for unknown remote stream 2" + ); + } + + #[test] + async fn packed_u8_admission_expands_offsets_and_acks_after_bus_acceptance() { + let (frames, _frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_server_with_capacity(frames, 1); + let value = session.decode(stream_reference(1)).await.unwrap(); + let mut primary = take_primary(value); + + let admission = tokio::spawn({ + let session = session.clone(); + async move { + session + .admit_input_item(InputStreamItem { + stream_id: 1, + sequence: 0, + payload: Some(input_stream_item::Payload::PackedU8(vec![7, 8, 9])), + }) + .await + } + }); + tokio::task::yield_now().await; + assert!(!admission.is_finished()); + + for (offset, value) in [(0, 7), (1, 8), (2, 9)] { + let event = primary.recv().await.unwrap(); + assert_eq!(event.offset, offset); + assert_eq!( + event.payload, + LiveStreamEventPayload::Item(SchemaValue::U8(value)) + ); + } + assert_eq!( + admission.await.unwrap().unwrap(), + InputItemAdmission::Acknowledged(InputStreamAck { + stream_id: 1, + sequence: 0, + logical_item_count: 3, + }) + ); + assert_eq!( + session + .admit_input_item(InputStreamItem { + stream_id: 1, + sequence: 2, + payload: Some(input_stream_item::Payload::Value( + SchemaValue::U8(10).try_into().unwrap(), + )), + }) + .await + .unwrap_err(), + "input stream 1 expected sequence 3, got 2" + ); + session.cancel(); + } + + #[test] + async fn input_sequence_overflow_is_rejected_before_recursive_registration() { + let (frames, _frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_server_with_capacity(frames, 1); + let value = session.decode(stream_reference(1)).await.unwrap(); + let _primary = take_primary(value); + let route = session.imported_route(1).unwrap(); + *route.next_sequence.lock().await = u64::MAX; + + assert_eq!( + session + .admit_input_item(InputStreamItem { + stream_id: 1, + sequence: u64::MAX, + payload: Some(input_stream_item::Payload::Value(stream_reference(3))), + }) + .await + .unwrap_err(), + "input stream 1 sequence overflow" + ); + assert!( + !session + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .states + .contains_key(&3) + ); + assert!( + !session + .inner + .seen_remote_stream_ids + .lock() + .expect("live stream ID set mutex poisoned") + .contains(&3) + ); + session.cancel(); + } + + #[test] + async fn imported_stream_error_is_scoped_and_terminal() { + let (frames, _frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_client(frames); + let value = session.decode(stream_reference(2)).await.unwrap(); + let mut primary = take_primary(value); + + session + .route_response(invocation_response::Response::OutputError( + OutputStreamError { + stream_id: 2, + offset: 0, + details: "failed".to_string(), + }, + )) + .await + .unwrap(); + + assert_eq!( + primary.recv().await.unwrap().payload, + LiveStreamEventPayload::Error("failed".to_string()) + ); + assert!(!session.is_cancelled()); + } +} + +#[cfg(test)] +mod bus_lifecycle_tests { + use super::*; + use crate::durable_host::stream_transport::{ + LiveStreamEndpoint, input_stream_pair, output_stream_pair, + }; + use std::time::Duration; + use test_r::{test, timeout}; + use tokio_util::sync::CancellationToken; + + fn stream_id(value: &ProtoSchemaValue) -> u64 { + match value.value.as_ref() { + Some(proto_schema_value::Value::StreamReference(reference)) => reference.stream_id, + other => panic!("expected stream id, got {other:?}"), + } + } + + fn stream_reference(id: u64) -> ProtoSchemaValue { + ProtoSchemaValue { + value: Some(proto_schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: id }, + )), + } + } + + fn stream_source(capacity: usize) -> (LiveStreamPeer, SchemaValue) { + let cancellation = CancellationToken::new(); + let (peer, stream) = input_stream_pair(capacity, &cancellation).unwrap(); + (peer, SchemaValue::Stream(stream)) + } + + fn take_primary( + value: SchemaValue, + ) -> crate::durable_host::stream_bus::PrimaryLiveStreamSubscriber { + let SchemaValue::Stream(stream) = value else { + panic!("expected stream"); + }; + stream + .take_host_endpoint::() + .unwrap() + .activate() + } + + #[test] + #[timeout("2s")] + async fn dropping_one_imported_input_reader_cancels_only_that_stream() { + let (frames, mut frame_rx) = mpsc::channel(16); + let session = LiveValueSession::new_server(frames); + let decoded = session + .decode(ProtoSchemaValue { + value: Some(proto_schema_value::Value::TupleValue(TupleValue { + elements: vec![stream_reference(1), stream_reference(3)], + })), + }) + .await + .unwrap(); + let SchemaValue::Tuple { mut elements } = decoded else { + panic!("expected tuple"); + }; + let second = elements.pop().unwrap(); + let first = elements.pop().unwrap(); + let first_endpoint = match first { + SchemaValue::Stream(stream) => { + stream.take_host_endpoint::().unwrap() + } + _ => unreachable!(), + }; + let mut second_primary = take_primary(second); + + drop(first_endpoint); + let Some(invocation_response::Response::StreamCancel(cancel)) = + frame_rx.recv().await.unwrap().response + else { + panic!("expected first stream cancellation"); + }; + assert_eq!(cancel.stream_id, 1); + assert_eq!(cancel.offset, 0); + assert_eq!(cancel.role(), StreamCancelRole::InputConsumer); + assert!(!session.is_cancelled()); + + for sequence in 0..2 { + session + .route_request(invocation_request::Request::InputItem(InputStreamItem { + stream_id: 1, + sequence, + payload: Some(input_stream_item::Payload::Value( + SchemaValue::Bool(false).try_into().unwrap(), + )), + })) + .await + .unwrap(); + } + session + .route_request(invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 1, + offset: 2, + })) + .await + .unwrap(); + assert!( + !session + .inner + .imported + .lock() + .expect("live stream map mutex poisoned") + .cancelled + .contains_key(&1) + ); + assert!(frame_rx.try_recv().is_err()); + + session + .route_request(invocation_request::Request::InputItem(InputStreamItem { + stream_id: 3, + sequence: 0, + payload: Some(input_stream_item::Payload::Value( + SchemaValue::Bool(true).try_into().unwrap(), + )), + })) + .await + .unwrap(); + assert_eq!( + second_primary.recv().await.unwrap().payload, + LiveStreamEventPayload::Item(SchemaValue::Bool(true)) + ); + assert!(matches!( + frame_rx.recv().await.unwrap().response, + Some(invocation_response::Response::InputAck(InputStreamAck { + stream_id: 3, + sequence: 0, + logical_item_count: 1, + })) + )); + session + .route_request(invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 3, + offset: 1, + })) + .await + .unwrap(); + assert_eq!( + second_primary.recv().await.unwrap().payload, + LiveStreamEventPayload::End + ); + session.wait_idle().await; + } + + #[test] + async fn equal_remote_ids_in_independent_sessions_do_not_alias() { + let (first_frames, _first_frame_rx) = mpsc::channel(8); + let first = LiveValueSession::new_client(first_frames); + let (second_frames, _second_frame_rx) = mpsc::channel(8); + let second = LiveValueSession::new_client(second_frames); + let mut first_primary = take_primary(first.decode(stream_reference(2)).await.unwrap()); + let mut second_primary = take_primary(second.decode(stream_reference(2)).await.unwrap()); + + first + .route_response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 2, + offset: 0, + value: Some(SchemaValue::String("first".to_string()).try_into().unwrap()), + }, + )) + .await + .unwrap(); + second + .route_response(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: 2, + offset: 0, + value: Some( + SchemaValue::String("second".to_string()) + .try_into() + .unwrap(), + ), + }, + )) + .await + .unwrap(); + + assert_eq!( + first_primary.recv().await.unwrap().payload, + LiveStreamEventPayload::Item(SchemaValue::String("first".to_string())) + ); + assert_eq!( + second_primary.recv().await.unwrap().payload, + LiveStreamEventPayload::Item(SchemaValue::String("second".to_string())) + ); + first.cancel(); + second.cancel(); + } + + #[test] + async fn invocation_finish_rejects_open_streams_and_releases_readers() { + let (frames, _frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_client(frames); + let mut primary = take_primary(session.decode(stream_reference(2)).await.unwrap()); + + let error = session.finish_invocation().await.unwrap_err(); + + assert!(error.contains("open imported streams [2]")); + assert_eq!( + primary.recv().await.unwrap_err(), + crate::durable_host::stream_bus::LiveStreamReceiveError::Closed + ); + session.wait_idle().await; + } + + #[test] + async fn streams_discovered_in_items_get_independent_buses() { + let (outer_source, outer_stream) = stream_source(4); + let (nested_source, nested_stream) = stream_source(4); + let (frames, mut frame_rx) = mpsc::channel(8); + let session = LiveValueSession::new_server(frames); + let outer_id = stream_id(&session.encode(&outer_stream).unwrap()); + outer_source + .publisher + .publish_item(SchemaValue::Option { + inner: Some(Box::new(nested_stream)), + }) + .await + .unwrap(); + let Some(invocation_response::Response::OutputItem(item)) = + frame_rx.recv().await.unwrap().response + else { + panic!("expected outer item"); + }; + let proto_schema_value::Value::OptionValue(option) = item.value.unwrap().value.unwrap() + else { + panic!("expected nested option"); + }; + let nested_id = stream_id(option.inner.as_deref().unwrap()); + assert_ne!(outer_id, nested_id); + + nested_source + .publisher + .publish_item(SchemaValue::String("nested".to_string())) + .await + .unwrap(); + let Some(invocation_response::Response::OutputItem(item)) = + frame_rx.recv().await.unwrap().response + else { + panic!("expected nested item"); + }; + assert_eq!(item.stream_id, nested_id); + assert_eq!( + SchemaValue::try_from(item.value.unwrap()).unwrap(), + SchemaValue::String("nested".to_string()) + ); + session.cancel(); + } + + #[test] + async fn output_bus_error_becomes_one_stream_error_frame() { + let (source, stream) = stream_source(4); + let (frames, mut frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_server(frames); + let id = stream_id(&session.encode(&stream).unwrap()); + source + .publisher + .publish_error("failed".to_string()) + .await + .unwrap(); + + assert!(matches!( + frame_rx.recv().await.unwrap().response, + Some(invocation_response::Response::OutputError(OutputStreamError { + stream_id, + offset: 0, + details, + })) if stream_id == id && details == "failed" + )); + session.wait_idle().await; + } + + #[test] + async fn output_primary_loss_cancels_the_invocation_tracker() { + let cancellation = CancellationToken::new(); + let tracker = Arc::new(LiveStreamTracker::new(cancellation.clone(), 4)); + let (consumer, stream) = output_stream_pair(Some(tracker.clone()), 4).unwrap(); + assert_eq!(tracker.active.load(Ordering::Acquire), 1); + + drop(stream); + + assert!(cancellation.is_cancelled()); + tokio::time::timeout(Duration::from_secs(1), tracker.wait_for_sources()) + .await + .unwrap(); + drop(consumer); + } + + #[test] + async fn normal_output_drop_finishes_after_sending_stream_end() { + let tracker = Arc::new(LiveStreamTracker::new(CancellationToken::new(), 4)); + let (consumer, stream) = output_stream_pair(Some(tracker.clone()), 4).unwrap(); + let (frames, mut frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_server(frames); + let stream_id = stream_id(&session.encode(&SchemaValue::Stream(stream)).unwrap()); + + drop(consumer); + tracker.wait_for_sources().await; + + session + .finish_invocation() + .await + .expect("a normal guest stream drop must wait for its end event"); + assert!(matches!( + frame_rx.recv().await.unwrap().response, + Some(invocation_response::Response::OutputEnd(OutputStreamEnd { + stream_id: actual_stream_id, + offset: 0, + })) if actual_stream_id == stream_id + )); + } + + #[test] + async fn cancellation_releases_an_exported_bus_waiter() { + let (_source, stream) = stream_source(4); + let (frames, _frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_server(frames); + session.encode(&stream).unwrap(); + + session.cancel(); + + tokio::time::timeout(Duration::from_secs(1), session.wait_idle()) + .await + .unwrap(); + } + + #[test] + async fn cancellation_releases_an_exported_stream_blocked_on_a_frame() { + let (source, stream) = stream_source(1); + let (frames, _frame_rx) = mpsc::channel(1); + frames + .send(InvocationResponse { response: None }) + .await + .unwrap(); + let session = LiveValueSession::new_server_with_capacity(frames, 1); + stream_id(&session.encode(&stream).unwrap()); + source + .publisher + .publish_item(SchemaValue::U32(42)) + .await + .unwrap(); + tokio::task::yield_now().await; + + session.cancel(); + + tokio::time::timeout(Duration::from_secs(1), session.wait_idle()) + .await + .unwrap(); + } + + #[test] + async fn unknown_stream_frames_are_rejected() { + let (frames, _frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_client(frames); + assert_eq!( + session + .route_response(invocation_response::Response::OutputEnd(OutputStreamEnd { + stream_id: 2, + offset: 0, + })) + .await + .unwrap_err(), + "item for unknown remote stream 2" + ); + } +} diff --git a/golem-worker-executor/src/durable_host/stream_transport.rs b/golem-worker-executor/src/durable_host/stream_transport.rs new file mode 100644 index 0000000000..931fca7c03 --- /dev/null +++ b/golem-worker-executor/src/durable_host/stream_transport.rs @@ -0,0 +1,579 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::durable_host::schema_value_stream::StoreValueResolver; +use crate::durable_host::stream_bus::{ + LiveStreamEventPayload, LiveStreamPublishError, LiveStreamPublisher, LiveStreamReceiveError, + PrimaryLiveStreamSubscriber, ReservedPrimaryLiveStreamSubscriber, live_input_stream_bus, + live_output_stream_bus, +}; +use crate::workerctx::WorkerCtx; +use golem_schema::schema::wit::wire::SchemaValueTree; +use golem_schema::schema::wit::{decode_value_with, encode_value_with_streams}; +use golem_schema::schema::{SchemaValue, SchemaValueStream}; +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; +use wasmtime::StoreContextMut; +use wasmtime::component::{Destination, Source, StreamConsumer, StreamProducer, StreamResult}; + +/// Tracks source endpoints created by one live streaming invocation. The +/// invocation keeps its Store event loop running until every source has +/// published its terminal or its primary reader has been lost. +#[derive(Debug)] +pub(crate) struct LiveStreamTracker { + pub(super) active: AtomicUsize, + changed: Notify, + cancelled: CancellationToken, + capacity: usize, +} + +impl LiveStreamTracker { + pub(crate) fn new(cancelled: CancellationToken, capacity: usize) -> Self { + assert!(capacity > 0, "live stream bus capacity must be non-zero"); + Self { + active: AtomicUsize::new(0), + changed: Notify::new(), + cancelled, + capacity, + } + } + + fn add_source(&self) { + self.active.fetch_add(1, Ordering::AcqRel); + self.changed.notify_waiters(); + } + + fn source_finished(&self) { + let previous = self.active.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "live stream source count underflow"); + self.changed.notify_waiters(); + } + + pub(crate) async fn wait_for_sources(&self) { + loop { + let changed = self.changed.notified(); + if self.active.load(Ordering::Acquire) == 0 { + return; + } + changed.await; + } + } + + async fn cancelled(&self) { + self.cancelled.cancelled().await; + } + + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancelled.clone() + } + + pub(crate) fn capacity(&self) -> usize { + self.capacity + } +} + +#[derive(Debug)] +pub(super) struct SourceLifecycle { + pub(super) finished: AtomicBool, + finished_notify: Notify, + trackers: Mutex>>, + cancelled: CancellationToken, +} + +impl SourceLifecycle { + fn new(cancelled: CancellationToken) -> Self { + Self { + finished: AtomicBool::new(false), + finished_notify: Notify::new(), + trackers: Mutex::new(Vec::new()), + cancelled, + } + } + + fn attach(self: &Arc, tracker: Arc) { + let mut trackers = self + .trackers + .lock() + .expect("stream lifecycle mutex poisoned"); + if self.finished.load(Ordering::Acquire) + || trackers + .iter() + .any(|current| Arc::ptr_eq(current, &tracker)) + { + return; + } + tracker.add_source(); + trackers.push(tracker.clone()); + let lifecycle = self.clone(); + tokio::spawn(async move { + tokio::select! { + _ = tracker.cancelled() => { + lifecycle.cancelled.cancel(); + lifecycle.finish(); + } + _ = lifecycle.wait_finished() => {} + } + }); + } + + async fn wait_finished(&self) { + loop { + let finished = self.finished_notify.notified(); + if self.finished.load(Ordering::Acquire) { + return; + } + finished.await; + } + } + + pub(super) fn finish(&self) { + if self.finished.swap(true, Ordering::AcqRel) { + return; + } + self.finished_notify.notify_waiters(); + let trackers = std::mem::take( + &mut *self + .trackers + .lock() + .expect("stream lifecycle mutex poisoned"), + ); + for tracker in trackers { + tracker.source_finished(); + } + } +} + +pub(super) struct LiveStreamEndpoint { + primary: Option>, + publisher: LiveStreamPublisher, + lifecycle: Arc, +} + +impl LiveStreamEndpoint { + pub(super) fn attach(&self, tracker: Arc) { + self.lifecycle.attach(tracker); + } + + pub(super) fn lifecycle(&self) -> Arc { + self.lifecycle.clone() + } + + pub(super) fn publisher(&self) -> LiveStreamPublisher { + self.publisher.clone() + } + + pub(super) fn activate(mut self) -> PrimaryLiveStreamSubscriber { + self.primary + .take() + .expect("live stream primary subscriber already activated") + .activate() + } +} + +impl Drop for LiveStreamEndpoint { + fn drop(&mut self) { + if self.primary.is_some() { + self.lifecycle.finish(); + } + } +} + +#[derive(Clone)] +pub(super) struct LiveStreamPeer { + pub(super) publisher: LiveStreamPublisher, + pub(super) primary_dropped: Arc, + pub(super) lifecycle: Arc, +} + +pub(super) fn input_stream_pair( + capacity: usize, + invocation_cancellation: &CancellationToken, +) -> Result<(LiveStreamPeer, SchemaValueStream), String> { + let stream_cancellation = invocation_cancellation.child_token(); + let primary_dropped = Arc::new(Notify::new()); + let lifecycle = Arc::new(SourceLifecycle::new(stream_cancellation.clone())); + let (publisher, primary) = + live_input_stream_bus(capacity, stream_cancellation, primary_dropped.clone()) + .map_err(|error| format!("failed to create live input stream bus: {error:?}"))?; + let endpoint = LiveStreamEndpoint { + primary: Some(primary), + publisher: publisher.clone(), + lifecycle: lifecycle.clone(), + }; + Ok(( + LiveStreamPeer { + publisher, + primary_dropped, + lifecycle, + }, + SchemaValueStream::from_host_endpoint(endpoint), + )) +} + +pub(super) fn output_stream_pair( + tracker: Option>, + capacity: usize, +) -> Result<(LiveOutputConsumer, SchemaValueStream), String> { + let cancellation = tracker + .as_ref() + .map(|tracker| tracker.cancellation_token()) + .unwrap_or_default(); + let lifecycle = Arc::new(SourceLifecycle::new(cancellation.clone())); + if let Some(tracker) = tracker { + debug_assert_eq!(capacity, tracker.capacity()); + lifecycle.attach(tracker); + } + let (publisher, primary) = live_output_stream_bus(capacity, cancellation) + .map_err(|error| format!("failed to create live output stream bus: {error:?}"))?; + let endpoint = LiveStreamEndpoint { + primary: Some(primary), + publisher: publisher.clone(), + lifecycle: lifecycle.clone(), + }; + Ok(( + LiveOutputConsumer { + publisher, + lifecycle, + pending: None, + pending_failure: None, + terminal_requested: false, + }, + SchemaValueStream::from_host_endpoint(endpoint), + )) +} + +type PublicationFuture = + Pin> + Send + 'static>>; + +pub(super) struct LiveOutputConsumer { + publisher: LiveStreamPublisher, + lifecycle: Arc, + pending: Option, + pending_failure: Option, + terminal_requested: bool, +} + +impl LiveOutputConsumer { + fn begin_terminal_publication(&mut self) { + self.terminal_requested = true; + let publisher = self.publisher.clone(); + self.pending = Some(Box::pin(async move { publisher.publish_end().await })); + } + + fn poll_pending(&mut self, cx: &mut Context<'_>) -> Poll> { + let result = match self.pending.as_mut() { + Some(pending) => match pending.as_mut().poll(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(result) => result, + }, + None => return Poll::Ready(Ok(StreamResult::Completed)), + }; + self.pending = None; + match result { + Ok(_) => match self.pending_failure.take() { + Some(_) => { + self.lifecycle.finish(); + Poll::Ready(Ok(StreamResult::Dropped)) + } + None if self.terminal_requested => { + self.lifecycle.finish(); + Poll::Ready(Ok(StreamResult::Cancelled)) + } + None => Poll::Ready(Ok(StreamResult::Completed)), + }, + Err(LiveStreamPublishError::Closed) => { + self.lifecycle.finish(); + Poll::Ready(Ok(StreamResult::Dropped)) + } + Err(error) => { + self.lifecycle.finish(); + Poll::Ready(Err(wasmtime::Error::msg(format!( + "failed to publish live output stream event: {error:?}" + )))) + } + } + } +} + +impl Drop for LiveOutputConsumer { + fn drop(&mut self) { + let pending = self.pending.take(); + let terminal_requested = self.terminal_requested; + let publisher = self.publisher.clone(); + let lifecycle = self.lifecycle.clone(); + tokio::spawn(async move { + if let Some(pending) = pending { + let _ = pending.await; + } + if !terminal_requested { + let _ = publisher.publish_end().await; + } + lifecycle.finish(); + }); + } +} + +impl StreamConsumer for LiveOutputConsumer { + type Item = SchemaValueTree; + + fn poll_consume( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + mut store: StoreContextMut, + mut source: Source<'_, Self::Item>, + finish: bool, + ) -> Poll> { + if self.pending.is_some() { + return self.poll_pending(cx); + } + if finish { + self.begin_terminal_publication(); + return self.poll_pending(cx); + } + + let mut item = None; + source.read(&mut store, &mut item)?; + let Some(item) = item else { + return Poll::Ready(Ok(StreamResult::Completed)); + }; + + let decoded = { + let mut resolver = StoreValueResolver::new(&mut store); + decode_value_with(item, &mut resolver).map_err(|error| error.to_string()) + }; + let publisher = self.publisher.clone(); + match decoded { + Ok(value) => { + self.pending = Some(Box::pin(async move { publisher.publish_item(value).await })); + } + Err(error) => { + self.pending_failure = Some(error.clone()); + self.terminal_requested = true; + self.pending = Some(Box::pin( + async move { publisher.publish_error(error).await }, + )); + } + } + self.poll_pending(cx) + } +} + +type ReceiveFuture = Pin< + Box< + dyn Future< + Output = ( + PrimaryLiveStreamSubscriber, + Option< + Result< + crate::durable_host::stream_bus::LiveStreamEvent, + LiveStreamReceiveError, + >, + >, + ), + > + Send + + 'static, + >, +>; + +async fn receive_input_event( + mut subscriber: PrimaryLiveStreamSubscriber, + cancelled: CancellationToken, +) -> ( + PrimaryLiveStreamSubscriber, + Option< + Result< + crate::durable_host::stream_bus::LiveStreamEvent, + LiveStreamReceiveError, + >, + >, +) { + let event = tokio::select! { + event = subscriber.recv() => Some(event), + _ = cancelled.cancelled() => None, + }; + (subscriber, event) +} + +pub(super) struct LiveInputProducer { + subscriber: Option>, + pending: Option, + lifecycle: Arc, + finished: bool, +} + +impl LiveInputProducer { + pub(super) fn new(endpoint: LiveStreamEndpoint) -> Self { + let lifecycle = endpoint.lifecycle.clone(); + Self { + subscriber: Some(endpoint.activate()), + pending: None, + lifecycle, + finished: false, + } + } +} + +impl Drop for LiveInputProducer { + fn drop(&mut self) { + self.lifecycle.finish(); + } +} + +impl StreamProducer for LiveInputProducer { + type Item = SchemaValueTree; + type Buffer = Option; + + fn poll_produce<'a>( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + mut store: StoreContextMut<'a, Ctx>, + mut destination: Destination<'a, Self::Item, Self::Buffer>, + finish: bool, + ) -> Poll> { + if self.finished { + return Poll::Ready(Ok(StreamResult::Dropped)); + } + if finish { + self.finished = true; + self.pending = None; + self.subscriber = None; + self.lifecycle.finish(); + return Poll::Ready(Ok(StreamResult::Cancelled)); + } + + if self.pending.is_none() { + let subscriber = self + .subscriber + .take() + .expect("live input stream subscriber is missing"); + let cancelled = self.lifecycle.cancelled.clone(); + self.pending = Some(Box::pin(receive_input_event(subscriber, cancelled))); + } + let (subscriber, event) = match self.pending.as_mut().unwrap().as_mut().poll(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(result) => result, + }; + self.pending = None; + self.subscriber = Some(subscriber); + + match event { + Some(Ok(event)) => match event.payload { + LiveStreamEventPayload::Item(value) => { + let encoded = { + let mut resolver = StoreValueResolver::new(&mut store); + encode_value_with_streams(&value, &mut resolver) + .map_err(|error| wasmtime::Error::msg(error.to_string()))? + }; + destination.set_buffer(Some(encoded)); + Poll::Ready(Ok(StreamResult::Completed)) + } + LiveStreamEventPayload::End | LiveStreamEventPayload::Cancel(_) => { + self.finished = true; + self.lifecycle.finish(); + Poll::Ready(Ok(StreamResult::Dropped)) + } + LiveStreamEventPayload::Error(error) => { + self.finished = true; + self.lifecycle.finish(); + Poll::Ready(Err(wasmtime::Error::msg(error))) + } + }, + Some(Err(LiveStreamReceiveError::Closed)) => { + self.finished = true; + self.lifecycle.finish(); + Poll::Ready(Err(wasmtime::Error::msg( + "live input stream closed without a terminal event", + ))) + } + Some(Err(LiveStreamReceiveError::Lagged(missed))) => { + self.finished = true; + self.lifecycle.finish(); + Poll::Ready(Err(wasmtime::Error::msg(format!( + "live input stream lost {missed} events" + )))) + } + None => { + self.finished = true; + self.subscriber = None; + self.lifecycle.finish(); + Poll::Ready(Err(wasmtime::Error::msg( + "live streaming invocation was cancelled", + ))) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use test_r::{test, timeout}; + + #[test] + #[timeout("2s")] + async fn normal_output_finish_publishes_end_before_finishing_lifecycle() { + let tracker = Arc::new(LiveStreamTracker::new(CancellationToken::new(), 4)); + let (mut consumer, stream) = output_stream_pair(Some(tracker.clone()), 4).unwrap(); + let endpoint = stream.take_host_endpoint::().unwrap(); + let mut primary = endpoint.activate(); + + consumer.begin_terminal_publication(); + let result = std::future::poll_fn(|cx| consumer.poll_pending(cx)) + .await + .unwrap(); + + assert!(matches!(result, StreamResult::Cancelled)); + assert!(consumer.lifecycle.finished.load(Ordering::Acquire)); + assert!(matches!( + primary.recv().await.unwrap(), + crate::durable_host::stream_bus::LiveStreamEvent { + offset: 0, + payload: LiveStreamEventPayload::End, + } + )); + tracker.wait_for_sources().await; + } + + #[test] + #[timeout("2s")] + async fn invocation_cancellation_wakes_a_guest_blocked_on_input() { + let invocation_cancellation = CancellationToken::new(); + let session_cancellation = CancellationToken::new(); + let tracker = Arc::new(LiveStreamTracker::new(invocation_cancellation.clone(), 4)); + let (peer, stream) = input_stream_pair(4, &session_cancellation).unwrap(); + let endpoint = stream.take_host_endpoint::().unwrap(); + endpoint.attach(tracker.clone()); + let primary = endpoint.activate(); + let blocked = tokio::spawn(receive_input_event( + primary, + peer.lifecycle.cancelled.clone(), + )); + tokio::task::yield_now().await; + assert!(!blocked.is_finished()); + + invocation_cancellation.cancel(); + + let (primary, event) = blocked.await.unwrap(); + assert_eq!(event, None); + drop(primary); + peer.primary_dropped.notified().await; + tracker.wait_for_sources().await; + assert!(!session_cancellation.is_cancelled()); + } +} diff --git a/golem-worker-executor/src/durable_host/wasm_rpc/mod.rs b/golem-worker-executor/src/durable_host/wasm_rpc/mod.rs index 936bdbe201..20e8fcea65 100644 --- a/golem-worker-executor/src/durable_host/wasm_rpc/mod.rs +++ b/golem-worker-executor/src/durable_host/wasm_rpc/mod.rs @@ -60,10 +60,11 @@ use golem_common::schema::schema_value::SchemaValue; use golem_common::serialization::{deserialize, serialize}; use golem_common::tracing::TraceOrigin; use golem_schema::schema::wit::{ - EncodeError, decode_typed_rejecting_quota_with, decode_value_with, encode_value_with, + EncodeError, decode_typed_rejecting_quota_with, decode_value_with, encode_value_with_streams, }; use crate::durable_host::golem::agent::schema_value_tree_to_typed_constructor_parameters; +use crate::worker::invocation::method_uses_streams; use golem_schema::schema::wit::wire as core_wire; use std::any::Any; use std::fmt::{Debug, Formatter}; @@ -256,6 +257,19 @@ impl HostWasmRpc for DurableWorkerCtx { let mut logical_remote_agent_id = OwnedAgentId::new(self.owned_agent_id.environment_id, &logical_remote_agent_id); + if self.is_live_streaming_invocation() { + return construct_ephemeral_wasm_rpc_resource( + self, + logical_remote_agent_id, + logical_agent_id, + env, + config, + span, + remote_agent_type, + component_revision, + ); + } + let mut handle = CallHandle::::start( self, HostRequestGolemRpcCreate { @@ -306,6 +320,19 @@ impl HostWasmRpc for DurableWorkerCtx { ); } + if self.is_live_streaming_invocation() { + return construct_live_wasm_rpc_resource( + self, + remote_agent_id, + &env, + config, + span, + remote_agent_type, + component_revision, + ) + .await; + } + let handle = CallHandle::::start( self, HostRequestGolemRpcCreate { @@ -372,11 +399,52 @@ impl HostWasmRpc for DurableWorkerCtx { "golem::rpc::wasm-rpc::invoke-and-await", method_name, input, + true, )? { Ok(prepared) => prepared, Err(err) => return Ok(Err(err)), }; + if prepared.is_streaming() { + if !self.state.is_live() { + return Ok(Err(RpcError::ProtocolError( + "live streaming invocation cannot be replayed".to_string(), + ))); + } + let idempotency_key = IdempotencyKey::fresh(); + let remote_agent_id = invocation_target_agent_id( + &prepared.logical_remote_agent_id, + prepared.ephemeral_logical_agent_id.as_ref(), + &idempotency_key, + )?; + let metadata = invocation_metadata(&remote_agent_id, &idempotency_key); + if !prepared.is_ephemeral() { + ensure_rpc_target_activated(self, self_).await?; + } + let result = self + .rpc() + .invoke_and_await_streaming( + &remote_agent_id, + Some(idempotency_key), + prepared.method_name, + prepared.input_value, + self.created_by(), + &self.agent_id().clone(), + &prepared.env, + self.clone_as_inherited_stack(&prepared.connection_span_id), + prepared.config, + &self.agent_auth_ctx(), + ) + .await; + return match result { + Ok(value) => Ok(Ok(InvocationResultWithMetadata { + metadata, + result: schema_value_to_wire_output(&value, self)?, + })), + Err(error) => Ok(Err(error.into())), + }; + } + let begun = CallHandle::::begin( self, DurableFunctionType::WriteRemote, @@ -440,6 +508,7 @@ impl HostWasmRpc for DurableWorkerCtx { "golem::rpc::wasm-rpc::invoke", method_name, input, + false, )? { Ok(prepared) => prepared, Err(err) => return Ok(Err(err)), @@ -534,15 +603,6 @@ impl HostWasmRpc for DurableWorkerCtx { )); } - // Check the per-invocation RPC call limit before initiating the call. - self.state - .check_and_increment_rpc_call_count() - .map_err(wasmtime::Error::from)?; - - // Returns Err(WorkerMonthlyRpcCallBudgetExhausted) when exhausted, - // which maps to RetryDecision::TryStop — suspending the worker. - self.record_monthly_rpc_call()?; - // Resolve the method and lift the input before opening any durability. Failures here are // deterministic functions of the cached remote agent type and the guest payload, so they // are baked into the future's result and surfaced on the first `get` — without opening a @@ -584,6 +644,86 @@ impl HostWasmRpc for DurableWorkerCtx { }); } }; + let method = find_agent_method(&remote_agent_type, &method_name)?; + let streaming = method_uses_streams(&remote_agent_type, method, &input_value); + + // Account for the call only after deterministic schema and mode classification has + // succeeded. Rejected calls must not consume per-invocation or monthly RPC budget. + self.state + .check_and_increment_rpc_call_count() + .map_err(wasmtime::Error::from)?; + self.record_monthly_rpc_call()?; + + if streaming { + if !self.state.is_live() { + let oplog_index = self.state.oplog.current_oplog_index().await; + let idempotency_key = self.derive_idempotency_key(oplog_index); + let remote_agent_id = invocation_target_agent_id( + &logical_remote_agent_id, + ephemeral_logical_agent_id.as_ref(), + &idempotency_key, + )?; + let metadata = invocation_metadata(&remote_agent_id, &idempotency_key); + let span = create_invocation_span( + self, + &connection_span_id, + &method_name, + &idempotency_key, + ) + .await?; + let fut = self.table().push(FutureInvokeResultEntry { + payload: Box::new(FutureInvokeResultState::Baked { + result: Ok(Err(InternalRpcError::ProtocolError { + details: "live streaming invocation cannot be replayed".to_string(), + })), + span_id: span.span_id().clone(), + }), + child_pollables: Vec::new(), + drop_pending: false, + })?; + return Ok(AsyncInvocationWithMetadata { + future: fut, + metadata, + }); + } + + let idempotency_key = IdempotencyKey::fresh(); + let remote_agent_id = invocation_target_agent_id( + &logical_remote_agent_id, + ephemeral_logical_agent_id.as_ref(), + &idempotency_key, + )?; + let metadata = invocation_metadata(&remote_agent_id, &idempotency_key); + let span = + create_invocation_span(self, &connection_span_id, &method_name, &idempotency_key) + .await?; + if ephemeral_logical_agent_id.is_none() { + ensure_rpc_target_activated(self, this).await?; + } + let task = spawn_streaming_invoke_and_await_task( + self, + remote_agent_id, + idempotency_key, + method_name, + input_value, + env, + config, + span.span_id(), + ); + let fut = self.table().push(FutureInvokeResultEntry { + payload: Box::new(FutureInvokeResultState::Live { + task: Some(Arc::new(tokio::sync::Mutex::new(task))), + span_id: span.span_id().clone(), + cancel_token: tokio_util::sync::CancellationToken::new(), + }), + child_pollables: Vec::new(), + drop_pending: false, + })?; + return Ok(AsyncInvocationWithMetadata { + future: fut, + metadata, + }); + } // Open the single durable host call for this async RPC as a `WriteRemote` — the same // durable function type as the synchronous `invoke_and_await`. It is a two-step call: @@ -817,7 +957,15 @@ impl DurableWorkerCtx { // invocation. let input_value = decode_value_with(input, self) .map_err(|err| anyhow::anyhow!("Invalid RPC input: {err}"))?; - find_agent_method(&remote_agent_type, &method_name)?; + let method = find_agent_method(&remote_agent_type, &method_name)?; + method + .validate_input(&remote_agent_type.schema, &input_value) + .map_err(|error| anyhow::anyhow!("Invalid RPC input: {error}"))?; + if method_uses_streams(&remote_agent_type, method, &input_value) { + return Err(anyhow::anyhow!( + "live streams cannot be used in scheduled invocations" + )); + } let scheduled_at = chrono::DateTime::from_timestamp(datetime.seconds, datetime.nanoseconds) .ok_or_else(|| { anyhow::Error::from(WorkerExecutorError::runtime(format!( @@ -980,6 +1128,7 @@ struct PreparedRpcInvocation { config: Vec, method_name: String, input_value: SchemaValue, + streaming: bool, } impl PreparedRpcInvocation { @@ -1001,6 +1150,10 @@ impl PreparedRpcInvocation { fn is_ephemeral(&self) -> bool { self.ephemeral_logical_agent_id.is_some() } + + fn is_streaming(&self) -> bool { + self.streaming + } } fn prepare_rpc_invocation( @@ -1009,6 +1162,7 @@ fn prepare_rpc_invocation( host_function_name: &str, method_name: String, input: core_wire::SchemaValueTree, + streaming_allowed: bool, ) -> anyhow::Result> { ctx.check_read_only_allows(host_function_name) .map_err(wasmtime::Error::from)?; @@ -1035,16 +1189,23 @@ fn prepare_rpc_invocation( ) }; - ctx.state - .check_and_increment_rpc_call_count() - .map_err(wasmtime::Error::from)?; - ctx.record_monthly_rpc_call()?; - let input_value = match resolve_method_and_lift_input(&remote_agent_type, &method_name, input, ctx) { Ok(input_value) => input_value, Err(err) => return Ok(Err(err.into())), }; + let method = remote_agent_type + .methods + .iter() + .find(|method| method.name == method_name) + .expect("method existence was checked while decoding RPC input"); + let streaming = method_uses_streams(&remote_agent_type, method, &input_value); + + if streaming && !streaming_allowed { + return Ok(Err(RpcError::ProtocolError( + "live streams require invoke-and-await".to_string(), + ))); + } if ephemeral_logical_agent_id.is_none() && logical_remote_agent_id == own_agent_id { return Err(anyhow::anyhow!( @@ -1052,6 +1213,13 @@ fn prepare_rpc_invocation( )); } + // Account for the call only after deterministic method, input, and stream + // classification has succeeded and before any remote or durable work. + ctx.state + .check_and_increment_rpc_call_count() + .map_err(wasmtime::Error::from)?; + ctx.record_monthly_rpc_call()?; + Ok(Ok(PreparedRpcInvocation { logical_remote_agent_id, ephemeral_logical_agent_id, @@ -1060,6 +1228,7 @@ fn prepare_rpc_invocation( config, method_name, input_value, + streaming, })) } @@ -1331,20 +1500,21 @@ async fn finish_span_access( accessor: &Accessor>>, span_id: &SpanId, ) -> Result<(), WorkerExecutorError> { - let (is_live, worker, replay_state) = accessor.with(|mut access| { + let (is_live, is_unpersisted_execution, worker, replay_state) = accessor.with(|mut access| { let ctx = access.get(); ( ctx.state.is_live(), + ctx.is_unpersisted_execution(), ctx.public_state.worker(), ctx.state.replay_state.clone(), ) }); - if is_live { + if is_live && !is_unpersisted_execution { worker .add_to_oplog(OplogEntry::finish_span(span_id.clone())) .await; - } else { + } else if !is_live { crate::get_oplog_entry_owned!(replay_state, OplogEntry::FinishSpan)?; } @@ -1469,6 +1639,13 @@ impl HostFutureInvokeResultWithStore span_id: SpanId, cancel_token: tokio_util::sync::CancellationToken, }, + /// A non-durable live session. No call handle or replay data exists; `get` only awaits + /// the result head and lowers it into the current Store. + Live { + task: FutureInvokeTaskHandle, + span_id: SpanId, + cancel_token: tokio_util::sync::CancellationToken, + }, } let plan = accessor.with(|mut access| { @@ -1544,6 +1721,17 @@ impl HostFutureInvokeResultWithStore cancel_token: cancel_token.clone(), } } + FutureInvokeResultState::Live { + task, + span_id, + cancel_token, + } => GetPlan::Live { + task: task + .take() + .ok_or_else(|| anyhow::anyhow!("future-invoke-result already consumed"))?, + span_id: span_id.clone(), + cancel_token: cancel_token.clone(), + }, }) })?; @@ -1558,6 +1746,57 @@ impl HostFutureInvokeResultWithStore } result } + GetPlan::Live { + task, + span_id, + cancel_token, + } => { + let interrupt_signal = accessor.with(|mut access| { + let ctx = access.get(); + ctx.create_interrupt_signal() + }); + let task_result = { + let mut guard = task.lock().await; + tokio::select! { + biased; + _ = cancel_token.cancelled() => None, + interrupt_kind = interrupt_signal => { + drop(guard); + drop(task); + return Err(interrupt_kind.into()); + } + result = &mut *guard => Some(result), + } + }; + let result = match task_result { + Some(result) => accessor.with(|mut access| { + future_invoke_get_result_to_wire( + future_invoke_task_result_to_get_result(&result), + access.get(), + ) + }), + None => Ok(Err(RpcError::ProtocolError( + "Invocation cancelled".to_string(), + ))), + }; + finish_span_access(accessor, &span_id).await?; + accessor.with(|mut access| { + let ctx = access.get(); + let entry = ctx + .table() + .get_mut(&Resource::::new_borrow(this_rep))?; + let state = entry + .payload + .as_any_mut() + .downcast_mut::() + .unwrap(); + *state = FutureInvokeResultState::Consumed { + span_id: span_id.clone(), + }; + Ok::<_, anyhow::Error>(()) + })?; + result + } GetPlan::Active { mut handle, task, @@ -1579,11 +1818,20 @@ impl HostFutureInvokeResultWithStore let (response, delivery) = if handle.is_live() { let task = task.expect("a live future-invoke-result must own its background task"); + let interrupt_signal = accessor.with(|mut access| { + let ctx = access.get(); + ctx.create_interrupt_signal() + }); let task_result = { let mut guard = task.lock().await; tokio::select! { biased; _ = cancel_token.cancelled() => None, + interrupt_kind = interrupt_signal => { + drop(guard); + drop(task); + return Err(handle.trap(interrupt_kind)); + } result = &mut *guard => Some(result), } }; @@ -1652,9 +1900,17 @@ impl HostFutureInvokeResultWithStore InvocationFreshnessDisposition::MayExist, ) }); + let interrupt_signal = accessor.with(|mut access| { + let ctx = access.get(); + ctx.create_interrupt_signal() + }); let task_result = tokio::select! { biased; _ = cancel_token.cancelled() => None, + interrupt_kind = interrupt_signal => { + drop(task); + return Err(live.trap(interrupt_kind)); + } result = &mut task => Some(result), }; match task_result { @@ -1822,6 +2078,17 @@ impl HostFutureInvokeResultWithStore FutureInvokeResultState::Baked { span_id, .. } => DropPlan::FinishSpan { span_id: span_id.clone(), }, + FutureInvokeResultState::Live { + task, + span_id, + cancel_token, + } => { + task.take(); + cancel_token.cancel(); + DropPlan::FinishSpan { + span_id: span_id.clone(), + } + } FutureInvokeResultState::Cancelled { .. } | FutureInvokeResultState::Consumed { .. } => DropPlan::Nothing, }) @@ -1882,6 +2149,8 @@ impl HostFutureInvokeResult for DurableWorkerCtx { idempotency_key: IdempotencyKey, span_id: SpanId, }, + /// A non-durable live session only needs its invocation span closed. + FinishSpan { span_id: SpanId }, /// Nothing to cancel: a baked failure, or already cancelled / consumed. Nothing, } @@ -1924,35 +2193,58 @@ impl HostFutureInvokeResult for DurableWorkerCtx { CancelPlan::Nothing } }, + FutureInvokeResultState::Live { + task, + span_id, + cancel_token, + } => { + if task.take().is_some() { + let span_id = span_id.clone(); + *state = FutureInvokeResultState::Cancelled { + span_id: span_id.clone(), + }; + CancelPlan::FinishSpan { span_id } + } else { + cancel_token.cancel(); + CancelPlan::Nothing + } + } FutureInvokeResultState::Baked { .. } | FutureInvokeResultState::Cancelled { .. } | FutureInvokeResultState::Consumed { .. } => CancelPlan::Nothing, } }; - if let CancelPlan::Cancel { - handle, - remote_agent_id, - idempotency_key, - span_id, - } = plan - { - // Best-effort remote cancellation, only meaningful for a live call — on replay the - // recorded `Cancelled` is re-applied without re-issuing the side effect. - if handle.is_live() - && let Err(err) = self - .worker_proxy() - .cancel_invocation(&remote_agent_id, idempotency_key, &self.agent_auth_ctx()) + match plan { + CancelPlan::Cancel { + handle, + remote_agent_id, + idempotency_key, + span_id, + } => { + // Best-effort remote cancellation, only meaningful for a live call — on replay the + // recorded `Cancelled` is re-applied without re-issuing the side effect. + if handle.is_live() + && let Err(err) = self + .worker_proxy() + .cancel_invocation( + &remote_agent_id, + idempotency_key, + &self.agent_auth_ctx(), + ) + .await + { + tracing::info!(err=%err, "Best-effort cancel_invocation failed"); + } + + handle + .cancel(self, None) .await - { - tracing::info!(err=%err, "Best-effort cancel_invocation failed"); + .map_err(anyhow::Error::from)?; + self.finish_span(&span_id).await?; } - - handle - .cancel(self, None) - .await - .map_err(anyhow::Error::from)?; - self.finish_span(&span_id).await?; + CancelPlan::FinishSpan { span_id } => self.finish_span(&span_id).await?, + CancelPlan::Nothing => {} } Ok(()) @@ -2007,22 +2299,6 @@ impl HostCancellationToken for DurableWorkerCtx { } } -impl core_wire::Host for DurableWorkerCtx { - async fn parse_uuid( - &mut self, - uuid: String, - ) -> anyhow::Result> { - Ok(uuid::Uuid::parse_str(&uuid) - .map(|uuid| uuid.into()) - .map_err(|e| e.to_string())) - } - - async fn uuid_to_string(&mut self, uuid: core_wire::Uuid) -> anyhow::Result { - let uuid: uuid::Uuid = uuid.into(); - Ok(uuid.to_string()) - } -} - fn construct_ephemeral_wasm_rpc_resource( ctx: &mut DurableWorkerCtx, remote_agent_id: OwnedAgentId, @@ -2066,6 +2342,51 @@ fn invocation_target_agent_id( )) } +async fn construct_live_wasm_rpc_resource( + ctx: &mut DurableWorkerCtx, + remote_agent_id: AgentId, + env: &[(String, String)], + config: Vec, + span: Arc, + remote_agent_type: Arc, + remote_component_revision: ComponentRevision, +) -> anyhow::Result> { + let stack = ctx.clone_as_inherited_stack(span.span_id()); + let target_component = ctx + .component_service() + .get_metadata(remote_agent_id.component_id, None) + .await?; + let remote_agent_id = OwnedAgentId::new(target_component.environment_id, &remote_agent_id); + let demand = ctx + .rpc() + .create_demand( + &remote_agent_id, + ctx.created_by(), + ctx.agent_id(), + env, + stack, + config.clone(), + &ctx.agent_auth_ctx(), + ) + .await?; + let target_fingerprint = demand.fingerprint(); + Ok(ctx.table().push(WasmRpcEntry { + payload: Box::new(WasmRpcEntryPayload { + remote_agent_id, + ephemeral_logical_agent_id: None, + span_id: span.span_id().clone(), + target_activation: WasmRpcTargetActivation::Activated { + demand, + target_fingerprint, + env: env.to_vec(), + config, + }, + remote_agent_type, + remote_component_revision, + }), + })?) +} + pub async fn construct_wasm_rpc_resource( ctx: &mut DurableWorkerCtx, mut handle: CallHandle, @@ -2274,6 +2595,7 @@ struct TaskRetryParams { runtime_retry_policy_mutations: std::collections::BTreeMap>, retry_properties: RetryProperties, max_in_function_retry_delay: Duration, + is_unpersisted_execution: bool, worker: Arc>, retry_point: OplogIndex, execution_status: Arc>, @@ -2440,6 +2762,7 @@ fn spawn_rpc_task_with_retry( max_in_function_retry_delay: retry_params.max_in_function_retry_delay, current_retry_policy_state, retry_properties: retry_params.retry_properties, + is_unpersisted_execution: retry_params.is_unpersisted_execution, worker: retry_params.worker, }; crate::durable_host::durability::in_task_retry_loop( @@ -2494,6 +2817,7 @@ fn spawn_invoke_and_await_task( runtime_retry_policy_mutations: ctx.state.runtime_retry_policy_mutations.clone(), retry_properties, max_in_function_retry_delay: ctx.durable_execution_state().max_in_function_retry_delay, + is_unpersisted_execution: ctx.is_unpersisted_execution(), worker: ctx.public_state.worker(), retry_point, execution_status: ctx.execution_status.clone(), @@ -2517,6 +2841,39 @@ fn spawn_invoke_and_await_task( ) } +fn spawn_streaming_invoke_and_await_task( + ctx: &mut DurableWorkerCtx, + remote_agent_id: OwnedAgentId, + idempotency_key: IdempotencyKey, + method_name: String, + input: SchemaValue, + env: Vec<(String, String)>, + config: Vec, + span_id: &SpanId, +) -> AbortOnDropJoinHandle { + let rpc = ctx.rpc(); + let created_by = ctx.created_by(); + let agent_id = ctx.agent_id().clone(); + let stack = ctx.clone_as_inherited_stack(span_id); + let auth_ctx = ctx.agent_auth_ctx(); + wasmtime_wasi::runtime::spawn(async move { + Ok(rpc + .invoke_and_await_streaming( + &remote_agent_id, + Some(idempotency_key), + method_name, + input, + created_by, + &agent_id, + &env, + stack, + config, + &auth_ctx, + ) + .await) + }) +} + pub struct WasmRpcEntryPayload { pub remote_agent_id: OwnedAgentId, pub ephemeral_logical_agent_id: Option, @@ -2601,7 +2958,7 @@ fn find_agent_method<'a>( agent_type: &'a AgentTypeSchema, method_name: &str, ) -> anyhow::Result<&'a AgentMethodSchema> { - agent_type + let method = agent_type .methods .iter() .find(|m| m.name == method_name) @@ -2610,7 +2967,8 @@ fn find_agent_method<'a>( "Method '{method_name}' not found on agent type '{}'", agent_type.type_name ) - }) + })?; + Ok(method) } /// Resolve and lift the guest-side input value tree into the schema-native @@ -2645,7 +3003,7 @@ fn resolve_method_and_lift_input( decode_value_with(input, resolver).map_err(|err| InternalRpcError::ProtocolError { details: format!("Invalid RPC input for method '{method_name}': {err}"), })?; - agent_type + let method = agent_type .methods .iter() .find(|m| m.name == method_name) @@ -2655,6 +3013,11 @@ fn resolve_method_and_lift_input( agent_type.type_name ), })?; + method + .validate_input(&agent_type.schema, &input_value) + .map_err(|error| InternalRpcError::ProtocolError { + details: format!("Invalid RPC input for method '{method_name}': {error}"), + })?; Ok(input_value) } @@ -2679,7 +3042,7 @@ fn schema_value_to_wire_output( ) -> Result, EncodeError> { match value { SchemaValue::Tuple { elements } if elements.is_empty() => Ok(None), - value => Ok(Some(encode_value_with(value, resolver)?)), + value => Ok(Some(encode_value_with_streams(value, resolver)?)), } } @@ -2762,6 +3125,14 @@ enum FutureInvokeResultState { span_id: SpanId, cancel_token: tokio_util::sync::CancellationToken, }, + /// A stream-bearing invocation is intentionally outside the oplog. The task owns the live + /// session until `get` takes it, while `cancel_token` lets a concurrent `cancel` stop an + /// in-flight `get` without serializing either the input or output value tree. + Live { + task: Option, + span_id: SpanId, + cancel_token: tokio_util::sync::CancellationToken, + }, /// Method resolution / input lifting failed deterministically before any host call was opened, /// so no `Start` / `End` is written for this future. `get` surfaces the baked error and finishes /// the span. Live and replay agree because the failure is a pure function of the cached remote @@ -2781,6 +3152,7 @@ impl Debug for FutureInvokeResultState { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { Self::Active { .. } => write!(f, "Active"), + Self::Live { .. } => write!(f, "Live"), Self::Baked { .. } => write!(f, "Baked"), Self::Cancelled { .. } => write!(f, "Cancelled"), Self::Consumed { .. } => write!(f, "Consumed"), diff --git a/golem-worker-executor/src/grpc/invocation.rs b/golem-worker-executor/src/grpc/invocation.rs index 80230bf6d4..4ab365bbfb 100644 --- a/golem-worker-executor/src/grpc/invocation.rs +++ b/golem-worker-executor/src/grpc/invocation.rs @@ -174,9 +174,7 @@ impl ProtobufInvocationDetails } } -impl ProtobufInvocationDetails - for golem_api_grpc::proto::golem::workerexecutor::v1::InvokeAgentRequest -{ +impl ProtobufInvocationDetails for golem_api_grpc::proto::golem::worker::InvocationStart { fn proto_agent_id(&self) -> &Option { &self.agent_id } @@ -222,7 +220,7 @@ mod tests { #[test] fn invoke_agent_request_decodes_creation_config() { - let request = golem_api_grpc::proto::golem::workerexecutor::v1::InvokeAgentRequest { + let request = golem_api_grpc::proto::golem::worker::InvocationStart { config: vec![golem_api_grpc::proto::golem::worker::AgentConfigEntryDto { path: vec!["database".into(), "port".into()], value: "5432".into(), @@ -238,7 +236,7 @@ mod tests { #[test] fn invoke_agent_request_rejects_invalid_creation_config_json() { - let request = golem_api_grpc::proto::golem::workerexecutor::v1::InvokeAgentRequest { + let request = golem_api_grpc::proto::golem::worker::InvocationStart { config: vec![golem_api_grpc::proto::golem::worker::AgentConfigEntryDto { path: vec!["database".into(), "port".into()], value: "not-json".into(), diff --git a/golem-worker-executor/src/grpc/invocation_session.rs b/golem-worker-executor/src/grpc/invocation_session.rs new file mode 100644 index 0000000000..9b4280493d --- /dev/null +++ b/golem-worker-executor/src/grpc/invocation_session.rs @@ -0,0 +1,970 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{WorkerExecutorImpl, extract_owned_agent_id}; +use crate::durable_host::stream_session::LiveValueSession; +use crate::grpc::invocation::{CanStartWorker, from_proto_invocation_context}; +use crate::services::{HasAll, HasComponentService, HasSchedulerService, UsesAllDeps}; +use crate::worker::Worker; +use crate::worker::invocation::validate_agent_method_invocation; +use crate::workerctx::WorkerCtx; +use chrono::{DateTime, Utc}; +use futures::{Stream, StreamExt}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem; +use golem_api_grpc::proto::golem::worker::v1::WorkerExecutionError; +use golem_api_grpc::proto::golem::worker::{ + InvocationAccepted, InvocationFailure, InvocationFailureKind, InvocationRejected, + InvocationRejectionReason, InvocationRequest, InvocationResponse, InvocationSessionCompletion, + InvocationSessionResult, InvocationStart, invocation_request, invocation_response, + invocation_session_completion, invocation_session_result, +}; +use golem_common::model::account::AccountId; +use golem_common::model::agent::{ + AgentMode, InvocationFreshnessDisposition, ParsedAgentId, Principal, +}; +use golem_common::model::component::ComponentRevision; +use golem_common::model::{ + AgentId, AgentInvocation, AgentInvocationOutput, AgentInvocationResult, IdempotencyKey, + InvocationStatus, ScheduledAction, +}; +use golem_common::schema::SchemaValue; +use golem_service_base::error::worker_executor::WorkerExecutorError; +use std::pin::Pin; +use std::sync::Arc; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{Request, Response, Status}; + +pub(super) type InvocationSessionStream = + Pin> + Send + 'static>>; + +pub(super) async fn invoke_agent_session< + Ctx: WorkerCtx, + Svcs: HasAll + UsesAllDeps + Send + Sync + 'static, +>( + executor: &WorkerExecutorImpl, + request: Request>, +) -> Result, Status> { + let inbound = request.into_inner(); + let (responses, receiver) = mpsc::channel(32); + let executor = (*executor).clone(); + tokio::spawn(async move { + executor.run_agent_session(inbound, responses).await; + }); + Ok(Response::new(Box::pin( + ReceiverStream::new(receiver).map(Ok), + ))) +} + +fn decode_invocation_freshness_disposition(value: i32) -> InvocationFreshnessDisposition { + if value == golem::worker::InvocationFreshnessDisposition::KnownFresh as i32 { + InvocationFreshnessDisposition::KnownFresh + } else { + InvocationFreshnessDisposition::MayExist + } +} + +fn publish_acceptance( + accepted: tokio::sync::oneshot::Sender>, + component_revision: Option, +) -> Result<(), WorkerExecutorError> { + accepted + .send(component_revision) + .map_err(|_| WorkerExecutorError::runtime("invocation session ended before acceptance")) +} + +impl + UsesAllDeps + Send + Sync + 'static> + WorkerExecutorImpl +{ + async fn invoke_agent_internal( + &self, + request: &InvocationStart, + method_parameters: Option, + cancellation: tokio_util::sync::CancellationToken, + accepted: tokio::sync::oneshot::Sender>, + ) -> Result { + Self::validate_auth_ctx(&request.auth_ctx)?; + + let freshness_disposition = + decode_invocation_freshness_disposition(request.freshness_disposition); + + let idempotency_key: Option = + request.idempotency_key.clone().map(|k| k.into()); + + if freshness_disposition == InvocationFreshnessDisposition::KnownFresh + && idempotency_key.is_none() + { + return Err(WorkerExecutorError::invalid_request( + "KnownFresh requires an idempotency key", + )); + } + + let mode = request.mode(); + + let ik = idempotency_key.unwrap_or(IdempotencyKey::fresh()); + let final_agent_id: AgentId = request + .agent_id + .clone() + .ok_or(WorkerExecutorError::invalid_request("agent_id not found"))? + .try_into() + .map_err(WorkerExecutorError::invalid_request)?; + + if matches!( + mode, + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup + ) { + if freshness_disposition == InvocationFreshnessDisposition::KnownFresh { + return Err(WorkerExecutorError::invalid_request( + "KnownFresh cannot be used for an invocation lookup", + )); + } + let inv_status = match self.get_or_create_pending_for_lookup(request).await? { + Some(worker) => match worker.lookup_invocation_result(&ik).await { + crate::model::LookupResult::Complete(Ok(_)) => InvocationStatus::Complete, + crate::model::LookupResult::Complete(Err(err)) => return Err(err), + crate::model::LookupResult::Pending => InvocationStatus::Pending, + crate::model::LookupResult::New | crate::model::LookupResult::Interrupted => { + InvocationStatus::Unknown + } + }, + None => InvocationStatus::Unknown, + }; + publish_acceptance(accepted, None)?; + return Ok(AgentInvocationOutput { + result: AgentInvocationResult::AgentInitialization, + consumed_fuel: None, + invocation_status: Some(inv_status), + component_revision: None, + agent_id: Some(final_agent_id), + idempotency_key: Some(ik), + oplog_index: None, + agent_fingerprint: None, + }); + } + + let method_name = + request + .method_name + .clone() + .ok_or(WorkerExecutorError::invalid_request( + "method_name is required for non-lookup invocations", + ))?; + + let method_parameters = method_parameters.ok_or(WorkerExecutorError::invalid_request( + "input is required for non-lookup invocations", + ))?; + + let schedule_at: Option> = request + .schedule_at + .as_ref() + .and_then(|ts| DateTime::from_timestamp(ts.seconds, ts.nanos as u32)); + + let account_id: AccountId = request + .component_owner_account_id + .ok_or(WorkerExecutorError::invalid_request("account_id not found"))? + .try_into() + .map_err(|e| { + WorkerExecutorError::invalid_request(format!("Invalid account id: {e}")) + })?; + + let owned_agent_id = + extract_owned_agent_id(request, |r| &r.agent_id, |r| &r.environment_id)?; + + Worker::::validate_invocation_freshness( + self, + &owned_agent_id, + &ik, + freshness_disposition, + ) + .await?; + + let principal: Principal = request + .principal + .clone() + .map(|p| p.try_into()) + .transpose() + .map_err(|e: String| { + WorkerExecutorError::invalid_request(format!("failed converting principal: {e}")) + })? + .unwrap_or_else(Principal::anonymous); + + let invocation_context = self + .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, + }; + + match mode { + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await => { + let worker = self + .get_or_create_pending_with_freshness(request, freshness_disposition) + .await?; + let status = worker.get_last_known_status().await; + let queued_manual_update_revision = status + .pending_invocations + .iter() + .rev() + .find_map(|invocation| invocation.manual_update_target_revision); + let pending_update_revision = queued_manual_update_revision.or_else(|| { + status + .pending_updates + .back() + .map(|update| update.target_revision) + }); + let current_component = self + .component_service() + .get_metadata( + owned_agent_id.component_id(), + Some(status.component_revision), + ) + .await?; + let (component, streaming) = if let Some(pending_revision) = pending_update_revision + && pending_revision != status.component_revision + && let Ok(pending_component) = self + .component_service() + .get_metadata(owned_agent_id.component_id(), Some(pending_revision)) + .await + && let Ok(parsed_agent_id) = ParsedAgentId::parse( + &owned_agent_id.agent_id.agent_id, + &pending_component.metadata, + ) + && let Ok(streaming) = validate_agent_method_invocation( + &pending_component.metadata, + Some(&parsed_agent_id), + &method_name, + &method_parameters, + ) { + (pending_component, streaming) + } else { + let parsed_agent_id = ParsedAgentId::parse( + &owned_agent_id.agent_id.agent_id, + ¤t_component.metadata, + ) + .map_err(WorkerExecutorError::invalid_request)?; + let streaming = validate_agent_method_invocation( + ¤t_component.metadata, + Some(&parsed_agent_id), + &method_name, + &method_parameters, + )?; + (current_component, streaming) + }; + let accepted_revision = + (worker.agent_mode() == AgentMode::Ephemeral).then_some(component.revision); + let mut invocation_output = if streaming { + let fingerprint = worker.get_initial_worker_metadata().fingerprint; + let invocation = worker + .clone() + .enqueue_live_streaming(invocation, cancellation) + .await?; + publish_acceptance(accepted, accepted_revision)?; + AgentInvocationOutput { + result: AgentInvocationResult::AgentMethod { + output: invocation.result().await?, + }, + consumed_fuel: None, + invocation_status: None, + component_revision: Some(component.revision), + agent_id: None, + idempotency_key: None, + oplog_index: None, + agent_fingerprint: Some(fingerprint), + } + } else { + publish_acceptance(accepted, accepted_revision)?; + worker.invoke_and_await(invocation).await? + }; + invocation_output.agent_id = Some(final_agent_id); + invocation_output.idempotency_key = Some(ik); + Ok(invocation_output) + } + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule => { + let existing_metadata = + if freshness_disposition == InvocationFreshnessDisposition::KnownFresh { + None + } else { + Worker::::get_latest_metadata(self, &owned_agent_id).await + }; + let component_revision = existing_metadata.as_ref().map(|metadata| { + let status = &metadata.last_known_status; + status + .pending_invocations + .iter() + .rev() + .find_map(|invocation| invocation.manual_update_target_revision) + .or_else(|| { + status + .pending_updates + .back() + .map(|update| update.target_revision) + }) + .unwrap_or(status.component_revision) + }); + let component = self + .component_service() + .get_metadata(owned_agent_id.component_id(), component_revision) + .await?; + let parsed_agent_id = + ParsedAgentId::parse(&owned_agent_id.agent_id.agent_id, &component.metadata) + .map_err(WorkerExecutorError::invalid_request)?; + let streaming = validate_agent_method_invocation( + &component.metadata, + Some(&parsed_agent_id), + &method_name, + &method_parameters, + )?; + if streaming { + return Err(WorkerExecutorError::invalid_request( + "live streams require an attached Await invocation session", + )); + } + + match schedule_at { + Some(scheduled_time) => { + let agent_mode = component + .metadata + .find_agent_type_by_name_ref(&parsed_agent_id.agent_type) + .map(|agent_type| agent_type.mode) + .ok_or_else(|| { + WorkerExecutorError::invalid_request( + "Scheduled invocation target is not a registered agent type", + ) + })?; + let action = if agent_mode == AgentMode::Ephemeral { + ScheduledAction::InvokeEphemeral { + account_id, + owned_agent_id, + invocation: Box::new(invocation), + component_revision: component.revision, + env: request.env().unwrap_or_default(), + config: request.config()?, + parent: request.parent(), + creation_principal: Box::new(worker_creation_principal), + } + } else { + let worker = self + .get_or_create_pending_with_freshness( + request, + freshness_disposition, + ) + .await?; + let target_worker_fingerprint = + worker.get_initial_worker_metadata().fingerprint; + ScheduledAction::Invoke { + account_id, + owned_agent_id, + invocation: Box::new(invocation), + target_worker_fingerprint, + } + }; + self.scheduler_service() + .schedule(scheduled_time, action) + .await; + publish_acceptance(accepted, Some(component.revision))?; + Ok(AgentInvocationOutput { + result: AgentInvocationResult::AgentInitialization, + consumed_fuel: None, + invocation_status: None, + component_revision: None, + agent_id: Some(final_agent_id), + idempotency_key: Some(ik), + oplog_index: None, + agent_fingerprint: None, + }) + } + None => { + let worker = self + .get_or_create_pending_with_freshness(request, freshness_disposition) + .await?; + let result = worker.clone().invoke(invocation).await?; + if let crate::worker::ResultOrSubscription::Finished(Err(err)) = &result { + return Err(err.clone()); + } + publish_acceptance(accepted, Some(component.revision))?; + match result { + crate::worker::ResultOrSubscription::Finished(Err(err)) => { + unreachable!( + "finished errors are handled before acceptance: {err}" + ); + } + crate::worker::ResultOrSubscription::Finished(Ok(_)) => {} + crate::worker::ResultOrSubscription::Pending(_) => { + Worker::start_if_needed(worker).await?; + } + } + Ok(AgentInvocationOutput { + result: AgentInvocationResult::AgentInitialization, + consumed_fuel: None, + invocation_status: None, + component_revision: None, + agent_id: Some(final_agent_id), + idempotency_key: Some(ik), + oplog_index: None, + agent_fingerprint: None, + }) + } + } + } + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup => { + unreachable!("Lookup mode handled above") + } + } + } + + async fn run_agent_session( + &self, + mut inbound: tonic::Streaming, + outward: mpsc::Sender, + ) { + let mut state = InvocationSessionState::default(); + let first = match inbound.message().await { + Ok(Some(request)) => request, + Ok(None) => { + send_unvalidated_rejection( + &outward, + InvocationRejectionReason::Protocol, + "invocation request transport closed before start".to_string(), + None, + None, + ) + .await; + return; + } + Err(error) => { + send_unvalidated_rejection( + &outward, + InvocationRejectionReason::Protocol, + error.to_string(), + None, + None, + ) + .await; + return; + } + }; + if let Err(error) = state.validate_trusted_request(&first) { + send_unvalidated_rejection( + &outward, + InvocationRejectionReason::Protocol, + error, + request_idempotency_key(&first), + request_agent_id(&first), + ) + .await; + return; + } + let first = first.request.expect("validated request has a payload"); + let start = match first { + invocation_request::Request::Start(start) => start, + invocation_request::Request::ResumeAttach(resume) => { + let rejection = InvocationResponse { + response: Some(invocation_response::Response::Rejected( + InvocationRejected { + reason: InvocationRejectionReason::ResumeUnsupported as i32, + error: "resume-attach is not supported by live sessions".to_string(), + idempotency_key: resume.idempotency_key, + agent_id: None, + component_revision: None, + }, + )), + }; + if state.validate_response(&rejection).is_ok() { + let _ = outward.send(rejection).await; + } + return; + } + _ => unreachable!("the session validator requires start or resume-attach first"), + }; + + let state = Arc::new(tokio::sync::Mutex::new(state)); + let (responses, mut response_rx) = mpsc::channel(32); + let response_state = state.clone(); + let outward_forwarder = outward.clone(); + let forwarder = tokio::spawn(async move { + while let Some(response) = response_rx.recv().await { + if let Err(error) = response_state.lock().await.validate_response(&response) { + tracing::error!(error, ?response, "Invalid invocation session response"); + return; + } + if outward_forwarder.send(response).await.is_err() { + return; + } + } + }); + + let session = LiveValueSession::new_server_with_capacity( + responses.clone(), + self.services + .config() + .limits + .live_stream_event_broadcast_capacity + .get(), + ); + let input = + if start.mode() == golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup { + Ok(None) + } else { + match start.input.clone() { + Some(input) => session.decode_start(input).await.map(Some), + None => Err("invocation start has no input".to_string()), + } + }; + let input = match input { + Ok(input) => input, + Err(error) => { + send_rejection( + &responses, + InvocationRejectionReason::Protocol, + error, + &start, + ) + .await; + session.cancel(); + drop(session); + drop(responses); + let _ = forwarder.await; + return; + } + }; + + let cancellation = tokio_util::sync::CancellationToken::new(); + let _cancel_on_drop = cancellation.clone().drop_guard(); + let (accepted_tx, mut accepted_rx) = tokio::sync::oneshot::channel(); + let invocation = self.invoke_agent_internal(&start, input, cancellation, accepted_tx); + tokio::pin!(invocation); + let mut early_output = None; + let accepted_revision = tokio::select! { + biased; + accepted = &mut accepted_rx => match accepted { + Ok(revision) => revision, + Err(_) => { + send_rejection( + &responses, + InvocationRejectionReason::Internal, + "invocation ended without reaching acceptance".to_string(), + &start, + ).await; + session.cancel(); + return; + } + }, + result = &mut invocation => { + if let Ok(revision) = accepted_rx.try_recv() { + early_output = Some(result); + revision + } else { + let (reason, error) = match result { + Ok(_) => ( + InvocationRejectionReason::Internal, + "invocation completed before acceptance".to_string(), + ), + Err(error) => ( + pre_acceptance_rejection_reason(&error), + error.to_string(), + ), + }; + send_rejection(&responses, reason, error, &start).await; + session.cancel(); + return; + } + } + request = inbound.message() => { + let error = match request { + Ok(Some(request)) => state + .lock() + .await + .validate_trusted_request(&request) + .unwrap_err(), + Ok(None) => "invocation request transport closed before acceptance".to_string(), + Err(error) => error.to_string(), + }; + send_rejection( + &responses, + InvocationRejectionReason::Protocol, + error, + &start, + ).await; + session.cancel(); + return; + } + }; + + if responses + .send(InvocationResponse { + response: Some(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: start.agent_id.clone(), + idempotency_key: start.idempotency_key.clone(), + component_revision: accepted_revision.map(|revision| revision.get()), + }, + )), + }) + .await + .is_err() + { + session.cancel(); + return; + } + + let output = match early_output { + Some(output) => output, + None => loop { + tokio::select! { + result = &mut invocation => break result, + request = inbound.message() => match request { + Ok(Some(request)) => { + if !route_live_request( + &session, + &responses, + &state, + request, + ).await { + return; + } + } + Ok(None) => { + fail_session_transport( + &session, + &responses, + "invocation request transport closed before the result".to_string(), + ) + .await; + return; + } + Err(error) => { + fail_session_transport(&session, &responses, error.to_string()).await; + return; + } + } + } + }, + }; + + let output = match output { + Ok(output) => output, + Err(error) => { + session.terminate_for_failure(&error.to_string()).await; + send_worker_failure(&responses, error).await; + return; + } + }; + let (result, output_stream_ids) = match &output.result { + AgentInvocationResult::AgentMethod { output } => match session.encode_pending(output) { + Ok((output, stream_ids)) => ( + Some(invocation_session_result::Result::MethodResult(output)), + stream_ids, + ), + Err(error) => { + session.terminate_for_failure(&error).await; + send_protocol_failure(&responses, error).await; + return; + } + }, + _ => ( + Some(invocation_session_result::Result::NoResult( + golem::common::Empty {}, + )), + Vec::new(), + ), + }; + if responses + .send(InvocationResponse { + response: Some(invocation_response::Response::Result( + InvocationSessionResult { + result, + component_revision: output + .component_revision + .map(|revision| revision.get()), + agent_id: output.agent_id.map(Into::into), + idempotency_key: output.idempotency_key.map(Into::into), + fuel_consumed: output.consumed_fuel, + status: output.invocation_status.map(|status| { + golem_api_grpc::proto::golem::worker::InvocationStatus::from(status) + as i32 + }), + oplog_index: output.oplog_index.map(u64::from), + agent_fingerprint: output + .agent_fingerprint + .map(|fingerprint| fingerprint.0.into()), + }, + )), + }) + .await + .is_err() + { + session.cancel(); + return; + } + session.activate_exported_streams(&output_stream_ids); + + let idle = session.wait_idle(); + tokio::pin!(idle); + loop { + tokio::select! { + () = &mut idle => { + if let Err(details) = session.finish_invocation().await { + send_protocol_failure(&responses, details).await; + return; + } + let _ = responses.send(InvocationResponse { + response: Some(invocation_response::Response::Finished( + InvocationSessionCompletion { + outcome: Some(invocation_session_completion::Outcome::Success( + golem::common::Empty {}, + )), + }, + )), + }).await; + return; + } + request = inbound.message() => match request { + Ok(Some(request)) => { + if !route_live_request( + &session, + &responses, + &state, + request, + ).await { + return; + } + } + Ok(None) => { + fail_session_transport( + &session, + &responses, + "invocation request transport closed before completion".to_string(), + ) + .await; + return; + } + Err(error) => { + fail_session_transport(&session, &responses, error.to_string()).await; + return; + } + } + } + } + } +} + +fn request_idempotency_key( + request: &InvocationRequest, +) -> Option { + match request.request.as_ref() { + Some(invocation_request::Request::Start(start)) => start.idempotency_key.clone(), + Some(invocation_request::Request::ResumeAttach(resume)) => resume.idempotency_key.clone(), + _ => None, + } +} + +fn request_agent_id( + request: &InvocationRequest, +) -> Option { + match request.request.as_ref() { + Some(invocation_request::Request::Start(start)) => start.agent_id.clone(), + _ => None, + } +} + +fn pre_acceptance_rejection_reason(error: &WorkerExecutorError) -> InvocationRejectionReason { + match error { + WorkerExecutorError::InvalidRequest { .. } + | WorkerExecutorError::ParamTypeMismatch { .. } + | WorkerExecutorError::NoValueInMessage + | WorkerExecutorError::ValueMismatch { .. } => InvocationRejectionReason::Validation, + WorkerExecutorError::AgentNotFound { .. } + | WorkerExecutorError::ComponentNotFound { .. } + | WorkerExecutorError::PromiseNotFound { .. } => InvocationRejectionReason::NotFound, + WorkerExecutorError::InvalidAccount => InvocationRejectionReason::Unauthorized, + _ => InvocationRejectionReason::Internal, + } +} + +async fn send_unvalidated_rejection( + responses: &mpsc::Sender, + reason: InvocationRejectionReason, + error: String, + idempotency_key: Option, + agent_id: Option, +) { + let _ = responses + .send(InvocationResponse { + response: Some(invocation_response::Response::Rejected( + InvocationRejected { + reason: reason as i32, + error, + idempotency_key, + agent_id, + component_revision: None, + }, + )), + }) + .await; +} + +async fn send_rejection( + responses: &mpsc::Sender, + reason: InvocationRejectionReason, + error: String, + start: &InvocationStart, +) { + send_unvalidated_rejection( + responses, + reason, + error, + start.idempotency_key.clone(), + start.agent_id.clone(), + ) + .await; +} + +async fn send_failure( + responses: &mpsc::Sender, + kind: InvocationFailureKind, + code: &str, + message: String, + worker_error: Option, +) { + let _ = responses + .send(InvocationResponse { + response: Some(invocation_response::Response::Finished( + InvocationSessionCompletion { + outcome: Some(invocation_session_completion::Outcome::Failure( + InvocationFailure { + kind: kind as i32, + code: code.to_string(), + message, + worker_error, + }, + )), + }, + )), + }) + .await; +} + +async fn send_protocol_failure(responses: &mpsc::Sender, details: String) { + send_failure( + responses, + InvocationFailureKind::Protocol, + "protocol", + details, + None, + ) + .await; +} + +async fn send_worker_failure( + responses: &mpsc::Sender, + error: WorkerExecutorError, +) { + let message = error.to_string(); + send_failure( + responses, + InvocationFailureKind::Execution, + "worker-execution", + message, + Some(error.into()), + ) + .await; +} + +async fn fail_session_transport( + session: &LiveValueSession, + responses: &mpsc::Sender, + details: String, +) { + session.terminate_for_failure(&details).await; + send_failure( + responses, + InvocationFailureKind::Transport, + "transport", + details, + None, + ) + .await; +} + +async fn route_live_request( + session: &LiveValueSession, + responses: &mpsc::Sender, + state: &Arc>, + request: InvocationRequest, +) -> bool { + let request = match state + .lock() + .await + .validate_received_trusted_request(&request) + { + Ok(()) => request + .request + .expect("validated invocation request has a payload"), + Err(details) => { + session.terminate_for_failure(&details).await; + send_protocol_failure(responses, details).await; + return false; + } + }; + match session.route_request(request).await { + Ok(true) => true, + Ok(false) => { + let details = "unexpected message on invocation request stream".to_string(); + session.terminate_for_failure(&details).await; + send_protocol_failure(responses, details).await; + false + } + Err(details) => { + session.terminate_for_failure(&details).await; + send_protocol_failure(responses, details).await; + false + } + } +} + +#[cfg(test)] +mod freshness_tests { + use super::decode_invocation_freshness_disposition; + use golem_common::model::agent::InvocationFreshnessDisposition; + use test_r::test; + + #[test] + fn invocation_freshness_defaults_unknown_values_to_may_exist() { + assert_eq!( + decode_invocation_freshness_disposition(0), + InvocationFreshnessDisposition::MayExist + ); + assert_eq!( + decode_invocation_freshness_disposition(i32::MAX), + InvocationFreshnessDisposition::MayExist + ); + } + + #[test] + fn invocation_freshness_decodes_known_fresh_explicitly() { + assert_eq!( + decode_invocation_freshness_disposition( + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh + as i32 + ), + InvocationFreshnessDisposition::KnownFresh + ); + } +} diff --git a/golem-worker-executor/src/grpc/mod.rs b/golem-worker-executor/src/grpc/mod.rs index eab184ca95..26c6164972 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -13,6 +13,7 @@ // limitations under the License. mod invocation; +mod invocation_session; use crate::grpc::invocation::{CanStartWorker, from_proto_invocation_context}; use crate::model::event::InternalWorkerEvent; @@ -27,17 +28,15 @@ use crate::services::worker_activator::{ use crate::services::worker_event::WorkerEventReceiver; use crate::services::{ All, HasActiveWorkers, HasAll, HasComponentService, HasEvents, HasOplogService, - HasPromiseService, HasRunningWorkerEnumerationService, HasSchedulerService, - HasShardManagerService, HasShardService, HasWorkerEnumerationService, HasWorkerService, - UsesAllDeps, + HasPromiseService, HasRunningWorkerEnumerationService, HasShardManagerService, HasShardService, + HasWorkerEnumerationService, HasWorkerService, UsesAllDeps, }; use crate::worker::Worker; use crate::workerctx::WorkerCtx; -use chrono::{DateTime, Utc}; use futures::Stream; use futures::StreamExt; use golem_api_grpc::proto::golem; -use golem_api_grpc::proto::golem::worker::{Cursor, UpdateMode}; +use golem_api_grpc::proto::golem::worker::{Cursor, InvocationRequest, UpdateMode}; use golem_api_grpc::proto::golem::workerexecutor::v1::worker_executor_server::WorkerExecutor; use golem_api_grpc::proto::golem::workerexecutor::v1::{ ActivatePluginRequest, ActivatePluginResponse, CancelInvocationRequest, @@ -46,10 +45,9 @@ use golem_api_grpc::proto::golem::workerexecutor::v1::{ GetAgentWalletRequest, GetAgentWalletResponse, GetAgentWalletSuccess, GetFileContentsRequest, GetFileContentsResponse, GetFileSystemNodeRequest, GetFileSystemNodeResponse, GetOplogRequest, GetOplogResponse, GetRunningWorkersMetadataRequest, GetRunningWorkersMetadataResponse, - GetWorkersMetadataRequest, GetWorkersMetadataResponse, InvokeAgentRequest, InvokeAgentResponse, - ProcessOplogEntriesRequest, ProcessOplogEntriesResponse, RevertWorkerRequest, - RevertWorkerResponse, SearchOplogRequest, SearchOplogResponse, UpdateWorkerRequest, - UpdateWorkerResponse, process_oplog_entries_response, + GetWorkersMetadataRequest, GetWorkersMetadataResponse, ProcessOplogEntriesRequest, + ProcessOplogEntriesResponse, RevertWorkerRequest, RevertWorkerResponse, SearchOplogRequest, + SearchOplogResponse, UpdateWorkerRequest, UpdateWorkerResponse, process_oplog_entries_response, }; use golem_common::metrics::api::record_new_grpc_api_active_stream; use golem_common::model::account::AccountId; @@ -64,11 +62,9 @@ use golem_common::model::oplog::{OplogIndex, UpdateDescription}; use golem_common::model::protobuf::to_protobuf_resource_description; use golem_common::model::worker::{AgentConfigEntryDto, AgentMetadataDto, TypedAgentConfigEntry}; use golem_common::model::{ - AgentEvent, AgentFilter, AgentFingerprint, AgentId, AgentInvocation, AgentInvocationOutput, - AgentInvocationResult, AgentMetadata, AgentStatus, IdempotencyKey, InvocationStatus, - OwnedAgentId, PendingUpdateKind, ScanCursor, ScheduledAction, ShardId, Timestamp, + AgentEvent, AgentFilter, AgentFingerprint, AgentId, AgentInvocation, AgentMetadata, + AgentStatus, IdempotencyKey, OwnedAgentId, PendingUpdateKind, ScanCursor, ShardId, Timestamp, }; -use golem_common::schema::SchemaValue; use golem_common::{model as common_model, recorded_grpc_api_request}; use golem_service_base::error::worker_executor::*; use golem_service_base::grpc::{ @@ -82,7 +78,6 @@ use std::marker::PhantomData; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use tokio; use tokio::sync::broadcast::error::RecvError; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tonic::{Request, Response, Status}; @@ -118,14 +113,6 @@ impl + UsesAllDeps + Send + Sync + type ResponseResult = Result, Status>; type ResponseStream = WorkerEventStream; -fn decode_invocation_freshness_disposition(value: i32) -> InvocationFreshnessDisposition { - if value == golem::workerexecutor::v1::InvocationFreshnessDisposition::KnownFresh as i32 { - InvocationFreshnessDisposition::KnownFresh - } else { - InvocationFreshnessDisposition::MayExist - } -} - impl + UsesAllDeps + Send + Sync + 'static> WorkerExecutorImpl { @@ -1889,252 +1876,6 @@ impl + UsesAllDeps + Send + Sync + } } - async fn invoke_agent_internal( - &self, - request: InvokeAgentRequest, - ) -> Result<(Option, Option), WorkerExecutorError> { - Self::validate_auth_ctx(&request.auth_ctx)?; - - let freshness_disposition = - decode_invocation_freshness_disposition(request.freshness_disposition); - - let idempotency_key: Option = - request.idempotency_key.clone().map(|k| k.into()); - - if freshness_disposition == InvocationFreshnessDisposition::KnownFresh - && idempotency_key.is_none() - { - return Err(WorkerExecutorError::invalid_request( - "KnownFresh requires an idempotency key", - )); - } - - let mode = request.mode(); - - let ik = idempotency_key.unwrap_or(IdempotencyKey::fresh()); - let final_agent_id: AgentId = request - .agent_id - .clone() - .ok_or(WorkerExecutorError::invalid_request("agent_id not found"))? - .try_into() - .map_err(WorkerExecutorError::invalid_request)?; - - if matches!( - mode, - golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup - ) { - if freshness_disposition == InvocationFreshnessDisposition::KnownFresh { - return Err(WorkerExecutorError::invalid_request( - "KnownFresh cannot be used for an invocation lookup", - )); - } - let inv_status = match self.get_or_create_pending_for_lookup(&request).await? { - Some(worker) => match worker.lookup_invocation_result(&ik).await { - crate::model::LookupResult::Complete(Ok(_)) => InvocationStatus::Complete, - crate::model::LookupResult::Complete(Err(err)) => return Err(err), - crate::model::LookupResult::Pending => InvocationStatus::Pending, - crate::model::LookupResult::New | crate::model::LookupResult::Interrupted => { - InvocationStatus::Unknown - } - }, - None => InvocationStatus::Unknown, - }; - return Ok(( - Some(AgentInvocationOutput { - result: AgentInvocationResult::AgentInitialization, - consumed_fuel: None, - invocation_status: Some(inv_status), - component_revision: None, - agent_id: Some(final_agent_id), - idempotency_key: Some(ik), - oplog_index: None, - agent_fingerprint: None, - }), - None, - )); - } - - let method_name = - request - .method_name - .clone() - .ok_or(WorkerExecutorError::invalid_request( - "method_name is required for non-lookup invocations", - ))?; - - let method_parameters: SchemaValue = request - .method_parameters - .clone() - .ok_or(WorkerExecutorError::invalid_request( - "method_parameters is required for non-lookup invocations", - ))? - .try_into() - .map_err(|e| { - WorkerExecutorError::invalid_request(format!( - "failed converting method_parameters: {e}" - )) - })?; - - let schedule_at: Option> = request - .schedule_at - .and_then(|ts| DateTime::from_timestamp(ts.seconds, ts.nanos as u32)); - - let account_id: AccountId = request - .component_owner_account_id - .ok_or(WorkerExecutorError::invalid_request("account_id not found"))? - .try_into() - .map_err(|e| { - WorkerExecutorError::invalid_request(format!("Invalid account id: {e}")) - })?; - - let owned_agent_id = - extract_owned_agent_id(&request, |r| &r.agent_id, |r| &r.environment_id)?; - - Worker::::validate_invocation_freshness( - self, - &owned_agent_id, - &ik, - freshness_disposition, - ) - .await?; - - let principal: Principal = request - .principal - .clone() - .map(|p| p.try_into()) - .transpose() - .map_err(|e: String| { - WorkerExecutorError::invalid_request(format!("failed converting principal: {e}")) - })? - .unwrap_or_else(Principal::anonymous); - - let invocation_context = self - .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, - input: method_parameters, - invocation_context, - principal, - }; - - match mode { - golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await => { - // Use the `pending` variant so we do NOT start the wasmtime instance - // up front. `Worker::invoke_and_await` checks the read-only cache first; - // on a cache hit (`ResultOrSubscription::Finished`) it returns without - // loading the agent. The Pending path starts the instance lazily so - // queued invocations still get processed. - let worker = self - .get_or_create_pending_with_freshness(&request, freshness_disposition) - .await?; - let mut invocation_output = worker.invoke_and_await(invocation).await?; - invocation_output.agent_id = Some(final_agent_id); - invocation_output.idempotency_key = Some(ik); - Ok((Some(invocation_output), None)) - } - golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule => { - match schedule_at { - Some(scheduled_time) => { - let component = self - .component_service() - .get_metadata(owned_agent_id.component_id(), None) - .await?; - let parsed_agent_id = ParsedAgentId::parse( - &owned_agent_id.agent_id.agent_id, - &component.metadata, - ) - .map_err(WorkerExecutorError::invalid_request)?; - let agent_mode = component - .metadata - .find_agent_type_by_name_ref(&parsed_agent_id.agent_type) - .map(|agent_type| agent_type.mode) - .ok_or_else(|| { - WorkerExecutorError::invalid_request( - "Scheduled invocation target is not a registered agent type", - ) - })?; - let action = if agent_mode == AgentMode::Ephemeral { - ScheduledAction::InvokeEphemeral { - account_id, - owned_agent_id, - invocation: Box::new(invocation), - component_revision: component.revision, - env: request.env().unwrap_or_default(), - config: request.config()?, - parent: request.parent(), - creation_principal: Box::new(worker_creation_principal), - } - } else { - let worker = self - .get_or_create_pending_with_freshness( - &request, - freshness_disposition, - ) - .await?; - let target_worker_fingerprint = - worker.get_initial_worker_metadata().fingerprint; - ScheduledAction::Invoke { - account_id, - owned_agent_id, - invocation: Box::new(invocation), - target_worker_fingerprint, - } - }; - self.scheduler_service() - .schedule(scheduled_time, action) - .await; - Ok(( - Some(AgentInvocationOutput { - result: AgentInvocationResult::AgentInitialization, - consumed_fuel: None, - invocation_status: None, - component_revision: None, - agent_id: Some(final_agent_id), - idempotency_key: Some(ik), - oplog_index: None, - agent_fingerprint: None, - }), - None, - )) - } - None => { - let worker = self - .get_or_create_pending_with_freshness(&request, freshness_disposition) - .await?; - match worker.clone().invoke(invocation).await? { - crate::worker::ResultOrSubscription::Finished(Err(err)) => { - return Err(err); - } - crate::worker::ResultOrSubscription::Finished(Ok(_)) => {} - crate::worker::ResultOrSubscription::Pending(_) => { - Worker::start_if_needed(worker).await?; - } - } - Ok(( - Some(AgentInvocationOutput { - result: AgentInvocationResult::AgentInitialization, - consumed_fuel: None, - invocation_status: None, - component_revision: None, - agent_id: Some(final_agent_id), - idempotency_key: Some(ik), - oplog_index: None, - agent_fingerprint: None, - }), - None, - )) - } - } - } - golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup => { - unreachable!("Lookup mode handled above") - } - } - } - async fn process_oplog_entries_internal( &self, request: ProcessOplogEntriesRequest, @@ -2284,7 +2025,16 @@ impl + UsesAllDeps + Send + Sync + agent_id: Some(metadata.agent_id.into()), environment_id: Some(metadata.environment_id.into()), env: HashMap::from_iter(metadata.env.iter().cloned()), - config: metadata.config.into_iter().map(Into::into).collect(), + config: metadata + .config + .into_iter() + .map(TryInto::try_into) + .collect::>() + .map_err(|error| { + WorkerExecutorError::unknown(format!( + "failed converting agent configuration: {error}" + )) + })?, created_by: Some(metadata.created_by.into()), component_revision: latest_status.component_revision.into(), status: Into::::into(latest_status.status).into(), @@ -3167,85 +2917,13 @@ impl + UsesAllDeps + Send + Sync + } } - async fn invoke_agent( - &self, - request: Request, - ) -> Result, Status> { - let request = request.into_inner(); - let record = recorded_grpc_api_request!( - "invoke_agent", - agent_id = proto_agent_id_string(&request.agent_id), - method_name = request.method_name.clone(), - idempotency_key = proto_idempotency_key_string(&request.idempotency_key), - ); + type InvokeAgentSessionStream = invocation_session::InvocationSessionStream; - match self - .invoke_agent_internal(request) - .instrument(record.span.clone()) - .await - { - Ok((result, _status)) => { - let ( - result_value, - fuel_consumed, - component_revision, - invocation_status, - oplog_index, - agent_fingerprint, - agent_id, - idempotency_key, - ) = match result { - Some(output) => { - let value = match &output.result { - AgentInvocationResult::AgentMethod { output } => { - Some(output.clone().into()) - } - _ => None, - }; - let proto_status = output.invocation_status.map(|s| { - golem_api_grpc::proto::golem::worker::InvocationStatus::from(s) as i32 - }); - ( - value, - output.consumed_fuel, - output.component_revision.map(|r| r.get()), - proto_status, - output.oplog_index.map(u64::from), - output.agent_fingerprint.map(|fp| fp.0.into()), - output.agent_id.map(Into::into), - output.idempotency_key.map(Into::into), - ) - } - None => (None, None, None, None, None, None, None, None), - }; - record.succeed(Ok(Response::new(InvokeAgentResponse { - result: Some( - golem::workerexecutor::v1::invoke_agent_response::Result::Success( - golem::workerexecutor::v1::InvokeAgentSuccess { - result: result_value, - fuel_consumed, - component_revision, - status: invocation_status, - oplog_index, - agent_fingerprint, - agent_id, - idempotency_key, - }, - ), - ), - }))) - } - Err(mut err) => record.fail( - Ok(Response::new(InvokeAgentResponse { - result: Some( - golem::workerexecutor::v1::invoke_agent_response::Result::Failure( - err.clone().into(), - ), - ), - })), - &mut err, - ), - } + async fn invoke_agent_session( + &self, + request: Request>, + ) -> ResponseResult { + invocation_session::invoke_agent_session(self, request).await } async fn process_oplog_entries( @@ -3356,33 +3034,3 @@ fn extract_owned_agent_id( Ok(OwnedAgentId::new(environment_id, &agent_id)) } - -#[cfg(test)] -mod freshness_tests { - use super::decode_invocation_freshness_disposition; - use golem_common::model::agent::InvocationFreshnessDisposition; - use test_r::test; - - #[test] - fn invocation_freshness_defaults_unknown_values_to_may_exist() { - assert_eq!( - decode_invocation_freshness_disposition(0), - InvocationFreshnessDisposition::MayExist - ); - assert_eq!( - decode_invocation_freshness_disposition(i32::MAX), - InvocationFreshnessDisposition::MayExist - ); - } - - #[test] - fn invocation_freshness_decodes_known_fresh_explicitly() { - assert_eq!( - decode_invocation_freshness_disposition( - golem_api_grpc::proto::golem::workerexecutor::v1::InvocationFreshnessDisposition::KnownFresh - as i32 - ), - InvocationFreshnessDisposition::KnownFresh - ); - } -} diff --git a/golem-worker-executor/src/lib.rs b/golem-worker-executor/src/lib.rs index 66eb1c352c..e289c887e8 100644 --- a/golem-worker-executor/src/lib.rs +++ b/golem-worker-executor/src/lib.rs @@ -343,9 +343,13 @@ pub trait Bootstrap { leak_sentinel: Arc<()>, ) -> anyhow::Result> { let worker_fork = Arc::new(DefaultWorkerFork::new( - Arc::new(RemoteInvocationRpc::new( + Arc::new(RemoteInvocationRpc::new_with_stream_capacity( worker_proxy.clone(), shard_service.clone(), + golem_config + .limits + .live_stream_event_broadcast_capacity + .get(), )), active_workers.clone(), engine.clone(), @@ -383,9 +387,13 @@ pub trait Bootstrap { )); let rpc = Arc::new(DirectWorkerInvocationRpc::new( - Arc::new(RemoteInvocationRpc::new( + Arc::new(RemoteInvocationRpc::new_with_stream_capacity( worker_proxy.clone(), shard_service.clone(), + golem_config + .limits + .live_stream_event_broadcast_capacity + .get(), )), direct_invocation_auth_service, active_workers.clone(), @@ -495,10 +503,10 @@ pub trait Bootstrap { &mut linker, DurableWorkerCtxView::durable_ctx_mut, )?; - golem_schema::schema::wit::wire::add_to_linker::<_, HasSelf>>( - &mut linker, - DurableWorkerCtxView::durable_ctx_mut, - )?; + golem_schema::schema::wit::wire::add_to_linker::< + _, + durable_host::schema_value_stream::CoreTypesHost, + >(&mut linker, DurableWorkerCtxView::durable_ctx_mut)?; Ok(linker) } } diff --git a/golem-worker-executor/src/services/golem_config.rs b/golem-worker-executor/src/services/golem_config.rs index 5e7e147013..94d58ce57d 100644 --- a/golem-worker-executor/src/services/golem_config.rs +++ b/golem-worker-executor/src/services/golem_config.rs @@ -32,6 +32,7 @@ use http::Uri; use serde::{Deserialize, Serialize}; use std::fmt::Write; use std::net::{Ipv4Addr, SocketAddrV4}; +use std::num::NonZeroUsize; use std::path::{Path, PathBuf}; use std::time::Duration; use tracing::warn; @@ -375,6 +376,7 @@ impl Default for GolemConfig { pub struct Limits { pub max_active_workers: usize, pub invocation_result_broadcast_capacity: usize, + pub live_stream_event_broadcast_capacity: NonZeroUsize, pub max_concurrent_streams: u32, pub event_broadcast_capacity: usize, pub event_history_size: usize, @@ -409,6 +411,11 @@ impl SafeDisplay for Limits { "invocation result broadcast capacity: {}", self.invocation_result_broadcast_capacity ); + let _ = writeln!( + &mut result, + "live stream event broadcast capacity: {}", + self.live_stream_event_broadcast_capacity + ); let _ = writeln!( &mut result, "max concurrent streams: {}", @@ -1662,6 +1669,7 @@ impl Default for Limits { Self { max_active_workers: 1024, invocation_result_broadcast_capacity: 100000, + live_stream_event_broadcast_capacity: NonZeroUsize::new(32).unwrap(), max_concurrent_streams: 1024, event_broadcast_capacity: 1024, event_history_size: 128, @@ -2066,3 +2074,40 @@ impl Default for QuotaServiceConfig { pub fn make_config_loader() -> ConfigLoader { ConfigLoader::new_with_examples(Path::new("config/worker-executor.toml")) } + +#[cfg(test)] +mod tests { + use super::Limits; + use golem_common::SafeDisplay; + use serde_json::Value; + use test_r::test; + + #[test] + fn live_stream_event_broadcast_capacity_defaults_to_32() { + let limits = Limits::default(); + + assert_eq!(limits.live_stream_event_broadcast_capacity.get(), 32); + let decoded: Limits = + serde_json::from_value(serde_json::to_value(&limits).unwrap()).unwrap(); + assert_eq!(decoded.live_stream_event_broadcast_capacity.get(), 32); + assert!( + limits + .to_safe_string() + .contains("live stream event broadcast capacity: 32") + ); + } + + #[test] + fn live_stream_event_broadcast_capacity_rejects_zero() { + let mut serialized = serde_json::to_value(Limits::default()).unwrap(); + let Value::Object(fields) = &mut serialized else { + panic!("limits must serialize as an object"); + }; + fields.insert( + "live_stream_event_broadcast_capacity".to_string(), + Value::from(0), + ); + + assert!(serde_json::from_value::(serialized).is_err()); + } +} diff --git a/golem-worker-executor/src/services/linear_memory.rs b/golem-worker-executor/src/services/linear_memory.rs index 2f074aa6ba..ffd90c04fb 100644 --- a/golem-worker-executor/src/services/linear_memory.rs +++ b/golem-worker-executor/src/services/linear_memory.rs @@ -42,7 +42,8 @@ struct Inner { startup_bytes_remaining: AtomicU64, pending_growth_prepaid: AtomicU64, growth_has_pending_grant: AtomicBool, - pending_growth_grants: Mutex>, + pending_growth_grants: Mutex>, + transient_growth_grants: Mutex>, retained_growth_grant: Arc>, reconciling: AtomicBool, replaying: AtomicBool, @@ -51,6 +52,12 @@ struct Inner { meter: AgentMemoryMeter, } +#[derive(Debug)] +struct PendingGrowthGrant { + grant: MemoryGrant, + transient: bool, +} + impl LinearMemoryTracker { pub(crate) fn new( bytes: u64, @@ -70,6 +77,7 @@ impl LinearMemoryTracker { pending_growth_prepaid: AtomicU64::new(0), growth_has_pending_grant: AtomicBool::new(false), pending_growth_grants: Mutex::new(Vec::new()), + transient_growth_grants: Mutex::new(Vec::new()), retained_growth_grant, reconciling: AtomicBool::new(true), replaying: AtomicBool::new(replaying), @@ -138,8 +146,16 @@ impl LinearMemoryTracker { .store(0, Ordering::Release); let pending_grants = std::mem::take(&mut *self.inner.pending_growth_grants.lock().unwrap()); let mut retained_grant = self.inner.retained_growth_grant.lock().unwrap(); - for grant in pending_grants { - retained_grant.merge(grant); + for pending in pending_grants { + if pending.transient { + self.inner + .transient_growth_grants + .lock() + .unwrap() + .push(pending.grant); + } else { + retained_grant.merge(pending.grant); + } } self.inner .growth_has_pending_grant @@ -154,13 +170,21 @@ impl LinearMemoryTracker { .inner .growth_has_pending_grant .swap(false, Ordering::AcqRel) - && let Some(grant) = self.inner.pending_growth_grants.lock().unwrap().pop() + && let Some(pending) = self.inner.pending_growth_grants.lock().unwrap().pop() { - self.inner - .retained_growth_grant - .lock() - .unwrap() - .merge(grant); + if pending.transient { + self.inner + .transient_growth_grants + .lock() + .unwrap() + .push(pending.grant); + } else { + self.inner + .retained_growth_grant + .lock() + .unwrap() + .merge(pending.grant); + } } let reconciling = self.inner.reconciling.load(Ordering::Acquire); let prepaid = self @@ -240,8 +264,20 @@ impl LinearMemoryTracker { } pub(crate) fn retain_growth_grant(&self, grant: MemoryGrant) { + self.retain_pending_growth_grant(grant, false); + } + + pub(crate) fn retain_transient_growth_grant(&self, grant: MemoryGrant) { + self.retain_pending_growth_grant(grant, true); + } + + fn retain_pending_growth_grant(&self, grant: MemoryGrant, transient: bool) { let _transition = self.inner.transitions.lock().unwrap(); - self.inner.pending_growth_grants.lock().unwrap().push(grant); + self.inner + .pending_growth_grants + .lock() + .unwrap() + .push(PendingGrowthGrant { grant, transient }); self.inner .growth_has_pending_grant .store(true, Ordering::Release); @@ -531,6 +567,41 @@ mod tests { ); } + #[test] + async fn transient_growth_grant_is_released_with_its_store_tracker() { + let now = Instant::now(); + let controller = Arc::new(AdmissionController::new( + Box::new(FixedProbe::new(100, 0)), + AdmissionPolicy { usable_ratio: 1.0 }, + )); + let retained_grant = Arc::new(Mutex::new( + controller.admit(40, &NoEvictionSource).await.unwrap(), + )); + let tracker = LinearMemoryTracker::new( + 40, + 40, + AgentMode::Durable, + false, + Arc::new(AtomicResourceEntry::new(0, 0, 0, 0, 0)), + retained_grant.clone(), + now, + ); + tracker.reconcile(40, now); + + let transient_grant = controller.admit(10, &NoEvictionSource).await.unwrap(); + tracker.retain_transient_growth_grant(transient_grant); + tracker.grow(10, now); + assert_eq!(controller.headroom_bytes(), 50); + + drop(tracker); + assert_eq!( + controller.headroom_bytes(), + 60, + "snapshot-only capacity must not be retained by the running worker" + ); + assert_eq!(retained_grant.lock().unwrap().bytes(), 40); + } + #[test] fn committed_growth_rechecks_a_concurrently_lowered_limit() { let now = Instant::now(); diff --git a/golem-worker-executor/src/services/oplog/plugin.rs b/golem-worker-executor/src/services/oplog/plugin.rs index 54e4298914..7302314c22 100644 --- a/golem-worker-executor/src/services/oplog/plugin.rs +++ b/golem-worker-executor/src/services/oplog/plugin.rs @@ -322,8 +322,13 @@ impl OplogProcessorPlugin for PerExecutorOplogProcessorPlugin>() + .map_err(|error| { + WorkerExecutorError::unknown(format!( + "failed converting agent configuration: {error}" + )) + })?, created_by: Some(worker_metadata.created_by.into()), component_revision: latest_status.component_revision.into(), status: Into::::into( diff --git a/golem-worker-executor/src/services/rpc.rs b/golem-worker-executor/src/services/rpc.rs index 378456cb37..386fe96e20 100644 --- a/golem-worker-executor/src/services/rpc.rs +++ b/golem-worker-executor/src/services/rpc.rs @@ -17,12 +17,13 @@ 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::LiveValueSession; use crate::durable_host::websocket::WebSocketConnectionPool; use crate::services::events::Events; use crate::services::oplog::plugin::OplogProcessorPlugin; use crate::services::resource_limits::ResourceLimits; use crate::services::shard::ShardService; -use crate::services::worker_proxy::{WorkerProxy, WorkerProxyError}; +use crate::services::worker_proxy::{InvocationResponseStream, WorkerProxy, WorkerProxyError}; use crate::services::{ HasActiveWorkers, HasAgentTypesService, HasBlobStoreService, HasCardService, HasComponentService, HasConfig, HasEvents, HasExtraDeps, HasFileLoader, HasHttpConnectionPool, @@ -36,13 +37,22 @@ use crate::services::{ worker_fork, }; use crate::worker::Worker; +use crate::worker::invocation::validate_agent_method_invocation; use crate::workerctx::WorkerCtx; use async_trait::async_trait; +use futures::StreamExt; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::worker::{ + InvocationFailure, InvocationFailureKind, InvocationRejected, InvocationRejectionReason, + InvocationRequest, InvocationStart, invocation_request, invocation_response, + invocation_session_completion, invocation_session_result, +}; use golem_common::model::account::AccountId; use golem_common::model::agent::{ - AgentInvocationMode, AgentPrincipal, InvocationFreshnessDisposition, Principal, + AgentInvocationMode, AgentPrincipal, InvocationFreshnessDisposition, ParsedAgentId, Principal, }; use golem_common::model::card::{AgentMethodName, AgentResourcePattern, AgentVerb}; +use golem_common::model::component::ComponentRevision; use golem_common::model::invocation_context::InvocationContextStack; use golem_common::model::oplog::types::SerializableRpcError; use golem_common::model::worker::AgentConfigEntryDto; @@ -54,11 +64,29 @@ use golem_service_base::error::worker_executor::WorkerExecutorError; use golem_service_base::model::auth::AuthCtx; use std::collections::HashMap; use std::fmt::{Display, Formatter}; +use std::future::Future; use std::sync::Arc; use tokio::runtime::Handle; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; use tracing::debug; use wasmtime_wasi_http::HttpConnectionPool; +async fn method_validation_revision( + freshness_disposition: InvocationFreshnessDisposition, + load_existing_revision: F, +) -> Option +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + if freshness_disposition == InvocationFreshnessDisposition::KnownFresh { + None + } else { + load_existing_revision().await + } +} + #[async_trait] pub trait Rpc: Send + Sync { async fn create_demand( @@ -87,6 +115,28 @@ pub trait Rpc: Send + Sync { auth_ctx: &AuthCtx, ) -> Result; + /// Executes an awaited invocation whose recursive value tree contains live + /// streams. This is a non-durable session: implementations must not route + /// the value through ordinary protobuf/oplog serialization. + async fn invoke_and_await_streaming( + &self, + _owned_agent_id: &OwnedAgentId, + _idempotency_key: Option, + _method_name: String, + _method_parameters: SchemaValue, + _self_created_by: AccountId, + _self_agent_id: &AgentId, + _self_env: &[(String, String)], + _self_stack: InvocationContextStack, + _config: Vec, + _auth_ctx: &AuthCtx, + ) -> Result { + Err(RpcError::ProtocolError { + details: "live streaming invocation is not supported by this RPC implementation" + .to_string(), + }) + } + async fn invoke( &self, owned_agent_id: &OwnedAgentId, @@ -103,6 +153,37 @@ pub trait Rpc: Send + Sync { ) -> Result<(), RpcError>; } +struct RemoteLiveRequestGuard { + session: LiveValueSession, + requests: Option>, +} + +impl RemoteLiveRequestGuard { + fn new(session: LiveValueSession, requests: mpsc::Sender) -> Self { + Self { + session, + requests: Some(requests), + } + } + + fn disarm(&mut self) { + self.requests = None; + } +} + +impl Drop for RemoteLiveRequestGuard { + fn drop(&mut self) { + let Some(requests) = self.requests.take() else { + return; + }; + if self.session.is_cancelled() { + return; + } + self.session.cancel(); + drop(requests); + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum RpcError { ProtocolError { details: String }, @@ -240,13 +321,27 @@ pub trait RpcDemand: Send + Sync { pub struct RemoteInvocationRpc { worker_proxy: Arc, _shard_service: Arc, + stream_capacity: usize, } impl RemoteInvocationRpc { pub fn new(worker_proxy: Arc, shard_service: Arc) -> Self { + Self::new_with_stream_capacity(worker_proxy, shard_service, 32) + } + + pub fn new_with_stream_capacity( + worker_proxy: Arc, + shard_service: Arc, + stream_capacity: usize, + ) -> Self { + assert!( + stream_capacity > 0, + "live stream bus capacity must be non-zero" + ); Self { worker_proxy, _shard_service: shard_service, + stream_capacity, } } } @@ -376,6 +471,165 @@ impl Rpc for RemoteInvocationRpc { } } + async fn invoke_and_await_streaming( + &self, + owned_agent_id: &OwnedAgentId, + idempotency_key: Option, + method_name: String, + method_parameters: SchemaValue, + _self_created_by: AccountId, + self_agent_id: &AgentId, + self_env: &[(String, String)], + self_stack: InvocationContextStack, + config: Vec, + auth_ctx: &AuthCtx, + ) -> Result { + let state = Arc::new(tokio::sync::Mutex::new(InvocationSessionState::default())); + let (requests, mut request_rx) = mpsc::channel(32); + let (wire_requests, receiver) = mpsc::channel(32); + let request_state = state.clone(); + tokio::spawn(async move { + while let Some(request) = request_rx.recv().await { + if request_state + .lock() + .await + .validate_trusted_request(&request) + .is_err() + { + return; + } + if wire_requests.send(request).await.is_err() { + return; + } + } + }); + let session = + LiveValueSession::new_client_with_capacity(requests.clone(), self.stream_capacity); + let (input, input_stream_ids) = session + .encode_pending(&method_parameters) + .map_err(|details| RpcError::ProtocolError { details })?; + requests + .send(InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(owned_agent_id.agent_id().into()), + method_name: Some(method_name), + input: Some(input), + idempotency_key: idempotency_key.map(Into::into), + context: Some(golem_api_grpc::proto::golem::worker::InvocationContext { + parent: Some(self_agent_id.clone().into()), + env: HashMap::from_iter(self_env.to_vec()), + tracing: Some(self_stack.into()), + }), + auth_ctx: Some(auth_ctx.clone().into()), + principal: Some(caller_agent_principal(self_agent_id).into()), + environment_id: Some(owned_agent_id.environment_id.into()), + config: config.into_iter().map(Into::into).collect(), + component_owner_account_id: None, + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + schedule_at: None, + freshness_disposition: golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + })), + }) + .await + .map_err(|_| RpcError::ProtocolError { + details: "live invocation request ended before start".to_string(), + })?; + let mut cancel_on_drop = RemoteLiveRequestGuard::new(session.clone(), requests.clone()); + let mut inbound = match self + .worker_proxy + .invoke_agent_session(Box::pin(ReceiverStream::new(receiver))) + .await + { + Ok(inbound) => inbound, + Err(error) => { + session.cancel(); + return Err(error.into()); + } + }; + + while let Some(response) = inbound.next().await { + let response = match response { + Ok(response) => response, + Err(error) => { + let details = error.to_string(); + fail_live_response_transport(&session, details).await; + return Err(error.into()); + } + }; + if let Err(details) = state.lock().await.validate_response(&response) { + fail_live_response_transport(&session, details.clone()).await; + return Err(RpcError::ProtocolError { details }); + } + match response.response { + Some(invocation_response::Response::Accepted(_)) => { + session.activate_exported_streams(&input_stream_ids); + } + Some(invocation_response::Response::Rejected(rejected)) => { + if let Err(error) = + confirm_terminal_response_is_last(&mut inbound, &state).await + { + session.fail(error.to_string()); + return Err(error); + } + let error = rpc_error_from_rejection(rejected); + session.fail(error.to_string()); + return Err(error); + } + Some(invocation_response::Response::Result(result)) => { + let output = match result.result { + Some(invocation_session_result::Result::MethodResult(output)) => { + match session.decode(output).await { + Ok(output) => output, + Err(details) => { + fail_live_response_transport(&session, details.clone()).await; + return Err(RpcError::ProtocolError { details }); + } + } + } + Some(invocation_session_result::Result::NoResult(_)) | None => { + let details = + "streaming agent invocation returned no method result".to_string(); + fail_live_response_transport(&session, details.clone()).await; + return Err(RpcError::ProtocolError { details }); + } + }; + spawn_live_response_router(session, inbound, requests, state); + cancel_on_drop.disarm(); + return Ok(output); + } + Some(invocation_response::Response::Finished(finished)) => { + if let Err(error) = + confirm_terminal_response_is_last(&mut inbound, &state).await + { + session.fail(error.to_string()); + return Err(error); + } + let error = rpc_error_from_invocation_finished(finished); + session.fail(error.to_string()); + return Err(error); + } + Some(response) => match session.route_response(response).await { + Ok(true) => {} + Ok(false) => { + let details = + "unexpected response before the invocation result".to_string(); + fail_live_response_transport(&session, details.clone()).await; + return Err(RpcError::ProtocolError { details }); + } + Err(details) => { + fail_live_response_transport(&session, details.clone()).await; + return Err(RpcError::ProtocolError { details }); + } + }, + None => unreachable!("response state validation rejects empty frames"), + } + } + let details = "invocation response ended before publishing a result".to_string(); + fail_live_response_transport(&session, details.clone()).await; + Err(RpcError::ProtocolError { details }) + } + async fn invoke( &self, owned_agent_id: &OwnedAgentId, @@ -415,6 +669,150 @@ impl Rpc for RemoteInvocationRpc { } } +fn spawn_live_response_router( + session: LiveValueSession, + mut inbound: InvocationResponseStream, + requests: mpsc::Sender, + state: Arc>, +) { + tokio::spawn(async move { + let _requests = requests; + while let Some(response) = inbound.next().await { + let response = match response { + Ok(response) => response, + Err(error) => { + fail_live_response_transport(&session, error.to_string()).await; + return; + } + }; + if let Err(details) = state.lock().await.validate_response(&response) { + fail_live_response_transport(&session, details).await; + return; + } + match response.response { + Some(invocation_response::Response::Finished(finished)) => { + if let Err(error) = + confirm_terminal_response_is_last(&mut inbound, &state).await + { + session.fail(error.to_string()); + return; + } + match finished.outcome { + Some(invocation_session_completion::Outcome::Success(_)) => { + if let Err(details) = session.finish_invocation().await { + session.fail(details); + } + } + Some(invocation_session_completion::Outcome::Failure(failure)) => { + session.fail(rpc_error_from_failure(failure).to_string()); + } + None => session.fail("invocation completion has no outcome".to_string()), + } + return; + } + Some(response) => match session.route_response(response).await { + Ok(true) => {} + Ok(false) => { + let details = "unexpected response after the invocation result".to_string(); + fail_live_response_transport(&session, details).await; + return; + } + Err(details) => { + fail_live_response_transport(&session, details).await; + return; + } + }, + None => unreachable!("response state validation rejects empty frames"), + } + } + if !state.lock().await.is_complete() { + fail_live_response_transport( + &session, + "invocation response ended before completion".to_string(), + ) + .await; + } + }); +} + +async fn confirm_terminal_response_is_last( + inbound: &mut InvocationResponseStream, + state: &Arc>, +) -> Result<(), RpcError> { + match inbound.next().await { + None => Ok(()), + Some(Err(error)) => Err(error.into()), + Some(Ok(response)) => { + let details = state.lock().await.validate_response(&response).unwrap_err(); + Err(RpcError::ProtocolError { details }) + } + } +} + +async fn fail_live_response_transport(session: &LiveValueSession, details: String) { + session.fail(details); +} + +fn rpc_error_from_rejection(rejected: InvocationRejected) -> RpcError { + match InvocationRejectionReason::try_from(rejected.reason) + .unwrap_or(InvocationRejectionReason::Internal) + { + InvocationRejectionReason::Unauthorized => RpcError::Denied { + details: rejected.error, + }, + InvocationRejectionReason::NotFound => RpcError::NotFound { + details: rejected.error, + }, + InvocationRejectionReason::Internal => RpcError::RemoteInternalError { + details: rejected.error, + }, + _ => RpcError::ProtocolError { + details: rejected.error, + }, + } +} + +fn rpc_error_from_failure(failure: InvocationFailure) -> RpcError { + if let Some(worker_error) = failure.worker_error { + return WorkerExecutorError::try_from(worker_error) + .map(Into::into) + .unwrap_or_else(|error| RpcError::RemoteInternalError { + details: format!("failed to decode worker execution error: {error}"), + }); + } + match InvocationFailureKind::try_from(failure.kind).unwrap_or(InvocationFailureKind::Internal) { + InvocationFailureKind::Protocol | InvocationFailureKind::Transport => { + RpcError::ProtocolError { + details: failure.message, + } + } + InvocationFailureKind::Execution | InvocationFailureKind::Internal => { + RpcError::RemoteInternalError { + details: failure.message, + } + } + InvocationFailureKind::Unspecified => RpcError::ProtocolError { + details: failure.message, + }, + } +} + +fn rpc_error_from_invocation_finished( + finished: golem_api_grpc::proto::golem::worker::InvocationSessionCompletion, +) -> RpcError { + match finished.outcome { + Some(invocation_session_completion::Outcome::Failure(failure)) => { + rpc_error_from_failure(failure) + } + Some(invocation_session_completion::Outcome::Success(_)) => RpcError::ProtocolError { + details: "invocation completed successfully before publishing a result".to_string(), + }, + None => RpcError::ProtocolError { + details: "invocation completion has no outcome".to_string(), + }, + } +} + fn caller_agent_principal(self_agent_id: &AgentId) -> Principal { Principal::Agent(AgentPrincipal { agent_id: self_agent_id.clone(), @@ -810,6 +1208,35 @@ impl DirectWorkerInvocationRpc { &owned_agent_id.agent_id, )) } + + async fn validate_method_invocation( + &self, + owned_agent_id: &OwnedAgentId, + method_name: &str, + method_parameters: &SchemaValue, + freshness_disposition: InvocationFreshnessDisposition, + ) -> Result { + let component_revision = method_validation_revision(freshness_disposition, || async { + Worker::::get_latest_metadata(self, owned_agent_id) + .await + .map(|metadata| metadata.last_known_status.component_revision) + }) + .await; + let component = self + .component_service() + .get_metadata(owned_agent_id.component_id(), component_revision) + .await?; + let parsed_agent_id = + ParsedAgentId::parse(&owned_agent_id.agent_id.agent_id, &component.metadata) + .map_err(|details| RpcError::ProtocolError { details })?; + validate_agent_method_invocation( + &component.metadata, + Some(&parsed_agent_id), + method_name, + method_parameters, + ) + .map_err(Into::into) + } } #[async_trait] @@ -918,6 +1345,20 @@ impl Rpc for DirectWorkerInvocationRpc { ) .await?; + if self + .validate_method_invocation( + owned_agent_id, + &method_name, + &method_parameters, + freshness_disposition, + ) + .await? + { + return Err(RpcError::ProtocolError { + details: "live streams require the attached streaming RPC".to_string(), + }); + } + let principal = caller_agent_principal(self_agent_id); let idempotency_key = idempotency_key.unwrap_or(IdempotencyKey::fresh()); Worker::::validate_invocation_freshness( @@ -977,6 +1418,84 @@ impl Rpc for DirectWorkerInvocationRpc { } } + async fn invoke_and_await_streaming( + &self, + owned_agent_id: &OwnedAgentId, + idempotency_key: Option, + method_name: String, + method_parameters: SchemaValue, + self_created_by: AccountId, + self_agent_id: &AgentId, + self_env: &[(String, String)], + self_stack: InvocationContextStack, + config: Vec, + auth_ctx: &AuthCtx, + ) -> Result { + let owned_agent_id = &self.canonicalize_owned_agent_id(owned_agent_id).await?; + if self + .shard_service() + .check_worker(&owned_agent_id.agent_id) + .is_err() + { + return self + .remote_rpc + .invoke_and_await_streaming( + owned_agent_id, + idempotency_key, + method_name, + method_parameters, + self_created_by, + self_agent_id, + self_env, + self_stack, + config, + auth_ctx, + ) + .await; + } + + self.direct_invocation_auth + .check( + self_created_by, + owned_agent_id, + AgentVerb::Invoke, + AgentResourcePattern::Method(AgentMethodName(method_name.clone())), + auth_ctx, + ) + .await?; + self.validate_method_invocation( + owned_agent_id, + &method_name, + &method_parameters, + InvocationFreshnessDisposition::MayExist, + ) + .await?; + let principal = caller_agent_principal(self_agent_id); + let worker = Worker::get_or_create_suspended( + self, + owned_agent_id, + Some(self_env.to_vec()), + config, + None, + Some(self_agent_id.clone()), + &self_stack, + principal.clone(), + ) + .await?; + let invocation = AgentInvocation::AgentMethod { + idempotency_key: idempotency_key.unwrap_or(IdempotencyKey::fresh()), + method_name, + input: method_parameters, + invocation_context: self_stack, + principal, + }; + let cancellation = tokio_util::sync::CancellationToken::new(); + worker + .invoke_live_streaming(invocation, cancellation) + .await + .map_err(Into::into) + } + async fn invoke( &self, owned_agent_id: &OwnedAgentId, @@ -1018,6 +1537,21 @@ impl Rpc for DirectWorkerInvocationRpc { ) .await?; + if self + .validate_method_invocation( + owned_agent_id, + &method_name, + &method_parameters, + freshness_disposition, + ) + .await? + { + return Err(RpcError::ProtocolError { + details: "live streams cannot be used in fire-and-forget invocations" + .to_string(), + }); + } + let principal = caller_agent_principal(self_agent_id); let idempotency_key = idempotency_key.unwrap_or(IdempotencyKey::fresh()); Worker::::validate_invocation_freshness( @@ -1075,3 +1609,78 @@ impl Rpc for DirectWorkerInvocationRpc { } } } + +#[cfg(test)] +mod protocol_tests { + use super::{ + RemoteLiveRequestGuard, RpcError, method_validation_revision, rpc_error_from_failure, + }; + use crate::durable_host::stream_session::LiveValueSession; + use golem_api_grpc::proto::golem::worker::{InvocationFailure, InvocationFailureKind}; + use golem_common::model::agent::InvocationFreshnessDisposition; + use golem_common::model::component::ComponentRevision; + use golem_service_base::error::worker_executor::WorkerExecutorError; + use std::cell::Cell; + use test_r::test; + use tokio::sync::mpsc; + + #[test] + async fn known_fresh_method_validation_uses_selected_revision_without_metadata_probe() { + let probed_existing_worker = Cell::new(false); + let revision = + method_validation_revision(InvocationFreshnessDisposition::KnownFresh, || async { + probed_existing_worker.set(true); + Some(ComponentRevision::INITIAL) + }) + .await; + + assert_eq!(revision, None); + assert!(!probed_existing_worker.get()); + } + + #[test] + async fn may_exist_method_validation_uses_existing_worker_revision() { + let probed_existing_worker = Cell::new(false); + let existing_revision = ComponentRevision::new(7).unwrap(); + let revision = + method_validation_revision(InvocationFreshnessDisposition::MayExist, || async { + probed_existing_worker.set(true); + Some(existing_revision) + }) + .await; + + assert_eq!(revision, Some(existing_revision)); + assert!(probed_existing_worker.get()); + } + + #[test] + fn typed_worker_failure_preserves_rpc_error_category() { + let worker_error = WorkerExecutorError::invalid_request("bad invocation"); + let error = rpc_error_from_failure(InvocationFailure { + kind: InvocationFailureKind::Execution as i32, + code: "worker-execution".to_string(), + message: worker_error.to_string(), + worker_error: Some(worker_error.into()), + }); + + assert_eq!( + error, + RpcError::ProtocolError { + details: "bad invocation".to_string(), + } + ); + } + + #[test] + async fn dropped_remote_request_cancels_the_session_and_closes_requests() { + let (frames, mut frame_rx) = mpsc::channel(4); + let session = LiveValueSession::new_client(frames.clone()); + let guard = RemoteLiveRequestGuard::new(session.clone(), frames); + + drop(guard); + + assert!(session.is_cancelled()); + drop(session); + assert!(frame_rx.recv().await.is_none()); + } +} diff --git a/golem-worker-executor/src/services/worker_proxy.rs b/golem-worker-executor/src/services/worker_proxy.rs index 477530789f..a00a6fa948 100644 --- a/golem-worker-executor/src/services/worker_proxy.rs +++ b/golem-worker-executor/src/services/worker_proxy.rs @@ -16,6 +16,7 @@ use super::golem_config::WorkerServiceGrpcConfig; use async_trait::async_trait; use chrono::{DateTime, Utc}; use desert_rust::BinaryCodec; +use futures::Stream; use golem_api_grpc::proto::golem::worker::v1::worker_service_client::WorkerServiceClient; use golem_api_grpc::proto::golem::worker::v1::{ AgentError, CancelInvocationRequest, CancelInvocationResponse, CompletePromiseRequest, @@ -27,7 +28,9 @@ use golem_api_grpc::proto::golem::worker::v1::{ invoke_agent_response, launch_new_worker_response, process_oplog_entries_response, resume_worker_response, revert_worker_response, update_worker_response, }; -use golem_api_grpc::proto::golem::worker::{CompleteParameters, UpdateMode}; +use golem_api_grpc::proto::golem::worker::{ + CompleteParameters, InvocationRequest, InvocationResponse, UpdateMode, +}; use golem_common::model::account::AccountId; use golem_common::model::agent::{AgentInvocationMode, InvocationFreshnessDisposition, Principal}; use golem_common::model::component::ComponentRevision; @@ -46,12 +49,39 @@ use golem_service_base::model::auth::AuthCtx; use std::collections::HashMap; use std::error::Error; use std::fmt::{Display, Formatter}; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; +use tonic::Status; use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use tonic_tracing_opentelemetry::middleware::client::OtelGrpcService; use tracing::debug; +pub type InvocationRequestStream = Pin + Send + 'static>>; +pub type InvocationResponseStream = + Pin> + Send + 'static>>; +type InvocationSessionCall<'a> = Pin< + Box< + dyn Future>, Status>> + + Send + + 'a, + >, +>; + +fn invoke_agent_session_once<'a>( + client: &'a mut WorkerServiceClient>, + request: Option, +) -> InvocationSessionCall<'a> { + match request { + Some(request) => Box::pin(client.invoke_agent_session(request)), + None => Box::pin(std::future::ready(Err(Status::aborted( + "invocation session request was already consumed", + )))), + } +} + #[async_trait] pub trait WorkerProxy: Send + Sync { async fn start( @@ -83,6 +113,17 @@ pub trait WorkerProxy: Send + Sync { auth_ctx: &AuthCtx, ) -> Result; + async fn invoke_agent_session( + &self, + _request: InvocationRequestStream, + ) -> Result { + Err(WorkerProxyError::InternalError( + WorkerExecutorError::invalid_request( + "invocation sessions are not supported by this worker proxy", + ), + )) + } + async fn update( &self, owned_agent_id: &OwnedAgentId, @@ -346,7 +387,11 @@ impl WorkerProxy for RemoteWorkerProxy { }); let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - method_parameters.into(); + method_parameters.try_into().map_err(|error| { + WorkerProxyError::BadRequest(vec![format!( + "method parameters cannot cross the remote worker boundary: {error}" + )]) + })?; let first_dispatch = AtomicBool::new(true); let response: InvokeAgentResponse = self @@ -458,6 +503,24 @@ impl WorkerProxy for RemoteWorkerProxy { } } + async fn invoke_agent_session( + &self, + request: InvocationRequestStream, + ) -> Result { + let request = Arc::new(std::sync::Mutex::new(Some(request))); + let response = self + .worker_service_client + .call_without_retry("invoke_agent_session", move |client| { + let request = request + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .take(); + invoke_agent_session_once(client, request) + }) + .await?; + Ok(Box::pin(response.into_inner())) + } + async fn update( &self, owned_agent_id: &OwnedAgentId, @@ -774,6 +837,9 @@ mod tests { #[tonic::async_trait] impl WorkerService for FlakyWorkerService { + type InvokeAgentSessionStream = + Pin> + Send + 'static>>; + unimplemented_rpc!( launch_new_worker, LaunchNewWorkerRequest, @@ -799,6 +865,13 @@ mod tests { ProcessOplogEntriesResponse ); + async fn invoke_agent_session( + &self, + _request: Request>, + ) -> Result, Status> { + Err(Status::unimplemented("invoke_agent_session")) + } + async fn invoke_agent( &self, request: Request, diff --git a/golem-worker-executor/src/worker/invocation.rs b/golem-worker-executor/src/worker/invocation.rs index 13dce176dd..5218cebd9d 100644 --- a/golem-worker-executor/src/worker/invocation.rs +++ b/golem-worker-executor/src/worker/invocation.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::durable_host::schema_value_stream::{StoreValueResolver, contains_stream}; +use crate::durable_host::stream_transport::LiveStreamTracker; use crate::metrics::wasm::{record_invocation, record_invocation_consumption}; use crate::model::TrapType; use crate::preview2::exports::golem::agent::guest as guest_exports; @@ -21,21 +23,27 @@ use crate::preview2::oplog_processor_plugin::exports::golem::api1_5_0::oplog_pro use crate::preview2::{golem_agent, golem_api_1_x}; use crate::workerctx::{PublicWorkerIo, WorkerCtx}; use futures::FutureExt; +use futures::channel::oneshot; use golem_common::model::agent::{AgentMode, ParsedAgentId}; use golem_common::model::component_metadata::ComponentMetadata; use golem_common::model::oplog::AgentError as OplogAgentError; use golem_common::model::{AgentInvocation, AgentInvocationResult, OplogIndex}; use golem_common::schema::SchemaValue; +#[cfg(test)] +use golem_common::schema::agent::InputSchema; use golem_common::schema::agent::wit::decode_agent_error_rejecting_quota_with; -use golem_common::schema::agent::{AgentTypeSchema, FieldSource, InputSchema}; +use golem_common::schema::agent::{AgentMethodSchema, AgentTypeSchema}; use golem_common::schema::graph::SchemaGraph; use golem_common::schema::schema_type::SchemaType; -use golem_common::schema::validation::value::{validate_record_fields, validate_value}; +use golem_common::schema::validation::value::validate_value; use golem_schema::schema::wit::wire as core_wire; -use golem_schema::schema::wit::{decode_value_with, encode_value_with}; +use golem_schema::schema::wit::{decode_value_with, encode_value_with, encode_value_with_streams}; use golem_service_base::error::worker_executor::{InterruptKind, WorkerExecutorError}; +use std::any::Any; +use std::sync::Arc; +use std::sync::Mutex; use tracing::{Instrument, Level, debug, span}; -use wasmtime::component::Accessor; +use wasmtime::component::{Accessor, AccessorTask}; use wasmtime::{AsContextMut, StoreContextMut}; /// Describes how an invocation is being executed with respect to the oplog. @@ -106,6 +114,220 @@ pub async fn invoke_observed_and_traced( } } +type LiveStreamingResponse = Result; +type LiveStreamingResponseSender = Arc>>>; + +fn publish_live_streaming_response( + response: &LiveStreamingResponseSender, + result: LiveStreamingResponse, +) -> bool { + response + .lock() + .expect("live streaming response mutex poisoned") + .take() + .is_some_and(|response| response.send(result).is_ok()) +} + +fn panic_payload_message(payload: &(dyn Any + Send)) -> String { + if let Some(message) = payload.downcast_ref::() { + message.clone() + } else if let Some(message) = payload.downcast_ref::<&str>() { + (*message).to_string() + } else { + "unknown panic payload".to_string() + } +} + +/// Runs one live, non-durable agent RPC whose value tree may contain recursive +/// streams. The result head is published as soon as `guest.invoke` returns; +/// the Store remains exclusively owned by this call until every stream sourced +/// by the guest reaches its terminal state or the downstream reader detaches. +pub async fn invoke_live_streaming_rpc( + lowered: LoweredInvocation, + store: &mut impl AsContextMut, + instance: &wasmtime::component::Instance, + response: oneshot::Sender>, + cancellation: tokio_util::sync::CancellationToken, +) -> Result<(), WorkerExecutorError> { + let mut store = store.as_context_mut(); + let response = Arc::new(Mutex::new(Some(response))); + let LoweredInvocation { + display_name, + read_only_method, + call, + } = lowered; + let guest = match load_agent_guest(&mut store, instance) { + Ok(guest) => guest, + Err(error) => { + publish_live_streaming_response(&response, Err(error.clone())); + return Err(error); + } + }; + if let Some(blocker) = store.data().snapshot_boundary_blocker() { + let error = WorkerExecutorError::runtime(format!( + "cannot start a live streaming invocation while {blocker}" + )); + publish_live_streaming_response(&response, Err(error.clone())); + return Err(error); + } + let LoweredCall::Invoke { + method_name, + input, + principal, + expected_output, + } = call + else { + let error = WorkerExecutorError::invalid_request( + "live streaming RPC is only supported for agent method invocations", + ); + publish_live_streaming_response(&response, Err(error.clone())); + return Err(error); + }; + let stream_capacity = store.data().durable_ctx().live_stream_event_capacity(); + let tracker = Arc::new(LiveStreamTracker::new(cancellation, stream_capacity)); + let mut tracker_installed = false; + let mut unpersisted_execution_active = false; + let mut read_only_active = false; + let mut invocation_started = false; + let response_for_call = response.clone(); + let tracker_for_call = tracker.clone(); + let validation_name = display_name.clone(); + + let call = std::panic::AssertUnwindSafe(async { + store + .data_mut() + .durable_ctx_mut() + .set_live_stream_tracker(tracker_for_call.clone()); + tracker_installed = true; + store + .data_mut() + .durable_ctx_mut() + .begin_unpersisted_streaming_invocation(); + unpersisted_execution_active = true; + + let input = { + let mut resolver = StoreValueResolver::new(&mut store); + encode_value_with_streams(&input, &mut resolver).map_err(|error| { + WorkerExecutorError::runtime(format!( + "Failed to encode live agent method input: {error}" + )) + })? + }; + + prepare_guest_call(&mut store, &display_name).await; + invocation_started = true; + store.data_mut().set_running(); + let _deadline = store.data().durable_ctx().arm_invocation_deadline(); + + if let Some(method_name) = &read_only_method { + store.data_mut().enter_read_only_mode(method_name.clone()); + read_only_active = true; + } + + run_guest_call_settled(&mut store, async move |accessor| { + match guest + .call_invoke(accessor, method_name, input, principal) + .await + { + Ok(Ok(output)) => { + let output = accessor.with(|mut access| match output { + None => Ok(SchemaValue::Tuple { + elements: Vec::new(), + }), + Some(tree) => { + let mut store = access.as_context_mut(); + let mut resolver = StoreValueResolver::new(&mut store); + decode_value_with(tree, &mut resolver) + .map_err(|error| wasmtime::Error::msg(error.to_string())) + } + })?; + validate_invoke_output(&validation_name, &expected_output, &output) + .map_err(|error| wasmtime::Error::msg(error.to_string()))?; + if publish_live_streaming_response(&response_for_call, Ok(output)) { + accessor + .spawn(WaitForLiveStreamSources { + tracker: tracker_for_call.clone(), + }) + .await; + } + Ok(true) + } + Ok(Err(error)) => { + let message = format!("agent method returned an error: {error:?}"); + publish_live_streaming_response( + &response_for_call, + Err(WorkerExecutorError::runtime(message)), + ); + Ok(false) + } + Err(error) => { + publish_live_streaming_response( + &response_for_call, + Err(WorkerExecutorError::runtime(error.to_string())), + ); + Err(error) + } + } + }) + .await + .and_then(|result| result) + .map_err(|error| WorkerExecutorError::runtime(error.to_string())) + }) + .catch_unwind() + .await; + + let (result, completed_successfully) = match call { + Ok(Ok(completed_successfully)) => (Ok(()), completed_successfully), + Ok(Err(error)) => (Err(error), false), + Err(payload) => ( + Err(WorkerExecutorError::runtime(format!( + "live streaming invocation panicked: {}", + panic_payload_message(payload.as_ref()) + ))), + false, + ), + }; + + if !completed_successfully { + tracker.cancellation_token().cancel(); + } + if read_only_active { + store.data_mut().exit_read_only_mode(); + } + if unpersisted_execution_active { + store + .data_mut() + .durable_ctx_mut() + .end_unpersisted_streaming_invocation_if_active(); + } + if tracker_installed { + store + .data_mut() + .durable_ctx_mut() + .clear_live_stream_tracker(); + } + if invocation_started { + let _ = finish_invocation_and_get_fuel_consumption(&mut store, &display_name).await?; + store.data().set_suspended(); + } + + if let Err(error) = &result { + publish_live_streaming_response(&response, Err(error.clone())); + } + result +} + +struct WaitForLiveStreamSources { + tracker: Arc, +} + +impl AccessorTask for WaitForLiveStreamSources { + async fn run(self, _accessor: &Accessor) -> wasmtime::Result<()> { + self.tracker.wait_for_sources().await; + Ok(()) + } +} + /// Invokes a worker and calls the appropriate hooks to observe the invocation async fn invoke_observed( lowered: LoweredInvocation, @@ -368,8 +590,8 @@ async fn dispatch_call( let consumed_fuel = finish_invocation_and_get_fuel_consumption(store, display_name).await?; match result { - Ok(Ok(maybe_output)) => { - let output = decode_invoke_output(store, maybe_output)?; + Ok(Ok(invoke_output)) => { + let output = decode_invoke_output(store, invoke_output)?; validate_invoke_output(display_name, &expected_output, &output)?; Ok(InvokeResult::Succeeded { consumed_fuel, @@ -527,26 +749,45 @@ fn invoke_result_from_agent_error( }) } -/// Decodes the `option` output of `invoke` into the -/// schema-native [`SchemaValue`] carried across the gRPC / oplog boundary. +/// Decodes the optional value returned by `invoke` into the schema-native +/// [`SchemaValue`] carried across the gRPC / oplog boundary. /// /// A `none` result (the declared `unit` output) is represented by the /// canonical empty tuple, matching the `unit` projection used on the caller /// side ([`schema_value_to_wire_output`](crate::durable_host::wasm_rpc)). fn decode_invoke_output( store: &mut StoreContextMut<'_, Ctx>, - maybe_output: Option, + output: Option, ) -> Result { - match maybe_output { + match output { // `none` is the declared `unit` output. None => Ok(SchemaValue::Tuple { elements: Vec::new(), }), - // The output is a guest-owned value tree, so any `quota-token` handles - // it carries are lifted into trusted snapshots (and consumed) here. - Some(tree) => decode_value_with(tree, store.data_mut().durable_ctx_mut()).map_err(|e| { - WorkerExecutorError::runtime(format!("Failed to decode agent method output: {e}")) - }), + // The output is crossing a durable/materializing boundary. Quota-token + // handles are lifted into trusted snapshots, while live stream handles + // are rejected and remain exclusive to `invoke_live_streaming_rpc`. + Some(tree) => { + let output = + decode_value_with(tree, store.data_mut().durable_ctx_mut()).map_err(|e| { + WorkerExecutorError::runtime(format!( + "Failed to decode agent method output: {e}" + )) + })?; + reject_stream_at_materializing_boundary(output) + } + } +} + +fn reject_stream_at_materializing_boundary( + output: SchemaValue, +) -> Result { + if contains_stream(&output) { + Err(WorkerExecutorError::runtime( + "Agent method output contains a live stream at a materializing invocation boundary", + )) + } else { + Ok(output) } } @@ -1000,12 +1241,7 @@ pub fn lower_invocation( })?; let read_only_method = method.read_only.is_some().then(|| method_name.clone()); - validate_schema_input_against_method_schema( - &input, - agent_type, - &method.input_schema, - &method_name, - )?; + validate_method_invocation(agent_type, method, &input, &method_name)?; let expected_output = Box::new(ExpectedInvokeOutput { graph: agent_type.schema.clone(), @@ -1082,51 +1318,49 @@ pub fn lower_invocation( } } -fn validate_schema_input_against_method_schema( - input: &SchemaValue, - agent_type: &AgentTypeSchema, - input_schema: &InputSchema, +pub fn validate_agent_method_invocation( + component_metadata: &ComponentMetadata, + agent_id: Option<&ParsedAgentId>, method_name: &str, -) -> Result<(), WorkerExecutorError> { - let SchemaValue::Record { fields } = input else { - return Err(WorkerExecutorError::invalid_request(format!( - "Method '{method_name}': expected input parameter record" - ))); - }; - - // Auto-injected fields (e.g. the principal) are supplied by the host to the - // guest export separately from the caller-provided input record, so they - // are excluded from both the parameter count and the value validation here. - let user_fields: Vec<_> = input_schema - .fields() + input: &SchemaValue, +) -> Result { + let agent_type = resolve_agent_type(component_metadata, agent_id)?; + let method = agent_type + .methods .iter() - .filter(|field| matches!(field.source, FieldSource::UserSupplied)) - .collect(); - if fields.len() != user_fields.len() { - return Err(WorkerExecutorError::invalid_request(format!( - "Method '{method_name}': expected {} parameters, got {}", - user_fields.len(), - fields.len() - ))); - } + .find(|method| method.name == method_name) + .ok_or_else(|| { + WorkerExecutorError::invalid_request(format!( + "Agent method '{method_name}' not found in agent type '{}'", + agent_type.type_name + )) + })?; + + validate_method_invocation(agent_type, method, input, method_name) +} - validate_record_fields( - &agent_type.schema, - user_fields - .iter() - .map(|field| (field.name.as_str(), &field.schema)), - fields, - ) - .map_err(|errors| { - WorkerExecutorError::invalid_request(format!( - "Method '{method_name}': invalid input parameter value: {}", - errors - .into_iter() - .map(|error| error.to_string()) - .collect::>() - .join("; ") - )) - }) +pub fn method_uses_streams( + agent_type: &AgentTypeSchema, + method: &AgentMethodSchema, + input: &SchemaValue, +) -> bool { + contains_stream(input) || method.uses_streams(&agent_type.schema) +} + +pub fn validate_method_invocation( + agent_type: &AgentTypeSchema, + method: &AgentMethodSchema, + input: &SchemaValue, + method_name: &str, +) -> Result { + method + .validate_input(&agent_type.schema, input) + .map_err(|error| { + WorkerExecutorError::invalid_request(format!( + "Method '{method_name}': invalid input parameter value: {error}" + )) + })?; + Ok(method_uses_streams(agent_type, method, input)) } /// Resolves the [`AgentTypeSchema`] an invocation targets: by name when an agent id @@ -1179,6 +1413,37 @@ mod tests { const AGENT_TYPE: &str = "test-agent"; const METHOD_NAME: &str = "do-work"; + #[test] + async fn live_streaming_response_is_published_exactly_once() { + let (sender, receiver) = oneshot::channel(); + let response = Arc::new(Mutex::new(Some(sender))); + + assert!(publish_live_streaming_response( + &response, + Ok(SchemaValue::Tuple { + elements: Vec::new(), + }) + )); + assert!(!publish_live_streaming_response( + &response, + Err(WorkerExecutorError::runtime("late failure")) + )); + + assert!(matches!( + receiver.await.unwrap(), + Ok(SchemaValue::Tuple { elements }) if elements.is_empty() + )); + } + + #[test] + fn live_streaming_panic_payload_is_reported() { + let owned: Box = Box::new("owned panic".to_string()); + let borrowed: Box = Box::new("borrowed panic"); + + assert_eq!(panic_payload_message(owned.as_ref()), "owned panic"); + assert_eq!(panic_payload_message(borrowed.as_ref()), "borrowed panic"); + } + /// Component metadata with one agent type whose `do-work` method takes two /// user-supplied parameters (`count: u32`, `label: string`) plus an /// auto-injected `principal` field. @@ -1200,6 +1465,10 @@ mod tests { http_endpoint: Vec::new(), read_only: None, }; + metadata_with_method(method) + } + + fn metadata_with_method(method: AgentMethodSchema) -> ComponentMetadata { let at = AgentTypeSchema { type_name: AgentTypeName(AGENT_TYPE.to_string()), description: String::new(), @@ -1271,7 +1540,7 @@ mod tests { panic!("non-record input must be rejected"); }; assert!( - err.to_string().contains("expected input parameter record"), + err.to_string().contains("expected record, found u32"), "unexpected error: {err}" ); } @@ -1289,7 +1558,7 @@ mod tests { panic!("arity mismatch must be rejected"); }; assert!( - err.to_string().contains("expected 2 parameters, got 1"), + err.to_string().contains("has 1 field(s), expected 2"), "unexpected error: {err}" ); } @@ -1312,6 +1581,89 @@ mod tests { ); } + #[test] + fn streaming_output_is_accepted_at_the_invocation_boundary() { + let metadata = metadata_with_method(AgentMethodSchema { + name: METHOD_NAME.to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(Vec::new()), + output_schema: OutputSchema::Single(Box::new(SchemaType::stream(Some( + SchemaType::u32(), + )))), + http_endpoint: Vec::new(), + read_only: None, + }); + let result = lower_invocation( + method_invocation(SchemaValue::Record { fields: Vec::new() }), + &metadata, + Some(&agent_id()), + ); + if let Err(error) = result { + panic!("unexpected error: {error}"); + } + } + + #[test] + fn streaming_method_is_classified_while_stream_free_method_is_not() { + let streaming = metadata_with_method(AgentMethodSchema { + name: METHOD_NAME.to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(Vec::new()), + output_schema: OutputSchema::Single(Box::new(SchemaType::stream(Some( + SchemaType::u32(), + )))), + http_endpoint: Vec::new(), + read_only: None, + }); + let empty_input = SchemaValue::Record { fields: Vec::new() }; + + assert!( + validate_agent_method_invocation( + &streaming, + Some(&agent_id()), + METHOD_NAME, + &empty_input, + ) + .unwrap() + ); + assert!( + !validate_agent_method_invocation( + &metadata_with_method(AgentMethodSchema { + name: METHOD_NAME.to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(Vec::new()), + output_schema: OutputSchema::Unit, + http_endpoint: Vec::new(), + read_only: None, + }), + Some(&agent_id()), + METHOD_NAME, + &empty_input, + ) + .unwrap() + ); + } + + #[test] + fn materializing_boundary_rejects_a_real_stream_handle() { + let stream = golem_common::schema::stream::SchemaValueStream::from_host_endpoint(()); + let output = SchemaValue::Record { + fields: vec![SchemaValue::Stream(stream)], + }; + + let error = reject_stream_at_materializing_boundary(output) + .expect_err("a live stream reaching materialization is a contract violation"); + assert!( + error + .to_string() + .contains("live stream at a materializing invocation boundary"), + "unexpected error: {error}" + ); + } + // --- validate_invoke_output --- fn unit_expected_output() -> ExpectedInvokeOutput { diff --git a/golem-worker-executor/src/worker/invocation_loop.rs b/golem-worker-executor/src/worker/invocation_loop.rs index 3ed46e9cdd..a771d6659e 100644 --- a/golem-worker-executor/src/worker/invocation_loop.rs +++ b/golem-worker-executor/src/worker/invocation_loop.rs @@ -19,12 +19,14 @@ use crate::services::linear_memory::LinearMemoryTracker; use crate::services::oplog::{CommitLevel, EphemeralOplog, OplogOps}; use crate::services::{HasEvents, HasOplog, HasWorker}; use crate::worker::invocation::{ - InvocationMode, InvokeResult, invoke_observed_and_traced, lower_invocation, + InvocationMode, InvokeResult, invoke_live_streaming_rpc, invoke_observed_and_traced, + lower_invocation, }; use crate::worker::status_checkpointer; use crate::worker::{ - FinalWorkerState, PendingWorkerInterrupt, QueuedWorkerInvocation, RetryDecision, RunningWorker, - Worker, WorkerCommand, WorkerInterruptState, WorkerTrace, + FinalWorkerState, PendingLiveInvocationDisposition, PendingWorkerInterrupt, + QueuedWorkerInvocation, RetryDecision, RunningWorker, Worker, WorkerCommand, + WorkerInterruptState, WorkerTrace, }; use crate::workerctx::{PublicWorkerIo, UpdateManagement, WorkerCtx}; use anyhow::anyhow; @@ -363,7 +365,12 @@ impl InvocationLoop { async fn stop_unloaded(&self, startup_failure: Option) { self.parent - .stop_internal(true, None, FinalWorkerState::Unloaded { startup_failure }) + .stop_internal( + true, + None, + FinalWorkerState::Unloaded { startup_failure }, + PendingLiveInvocationDisposition::Fail, + ) .await; } @@ -437,6 +444,7 @@ impl InvocationLoop { FinalWorkerState::Unloaded { startup_failure: Some(err), }, + PendingLiveInvocationDisposition::Fail, ) .await; CreateInstanceResult::Failed @@ -488,6 +496,7 @@ impl InvocationLoop { FinalWorkerState::Unloaded { startup_failure: Some(err), }, + PendingLiveInvocationDisposition::Fail, ) .await; Some(RetryDecision::None) // early return, we can't retry this @@ -594,7 +603,7 @@ impl InnerInvocationLoop<'_, Ctx> { break self.interrupt(interrupt).await; } - let message = self.active.write().await.pop_front(); + let message = self.pop_ready_internal_invocation().await; let result = if let Some(message) = message { self.internal_invocation(message).await @@ -666,6 +675,21 @@ impl InnerInvocationLoop<'_, Ctx> { } } + async fn pop_ready_internal_invocation(&self) -> Option { + let live_invocation_is_next = matches!( + self.active.read().await.front(), + Some(QueuedWorkerInvocation::LiveStreamingInvocation { .. }) + ); + if live_invocation_is_next { + let status = self.parent.get_non_detached_last_known_status().await; + if !status.pending_updates.is_empty() || !status.pending_invocations.is_empty() { + return None; + } + } + + self.active.write().await.pop_front() + } + /// Checks — before publishing `waiting_for_command = true`, which makes the /// worker eligible for idle eviction — that no Golem-spawned store task is /// still active. Every guest call drains its tail work before returning @@ -963,6 +987,7 @@ impl InnerInvocationLoop<'_, Ctx> { FinalWorkerState::Unloaded { startup_failure: Some(err), }, + PendingLiveInvocationDisposition::Fail, ) .await; CommandOutcome::BreakOuterLoop @@ -1040,10 +1065,66 @@ impl Invocation<'_, Ctx> { let _ = sender.send(Ok(())); CommandOutcome::Continue } + QueuedWorkerInvocation::LiveStreamingInvocation { + invocation, + sender, + cancellation, + } => { + self.invoke_live_streaming(*invocation, sender, cancellation) + .await + } QueuedWorkerInvocation::SaveSnapshot => self.save_snapshot().await, } } + async fn invoke_live_streaming( + &mut self, + invocation: AgentInvocation, + sender: oneshot::Sender>, + cancellation: tokio_util::sync::CancellationToken, + ) -> CommandOutcome { + if cancellation.is_cancelled() { + let _ = sender.send(Err(WorkerExecutorError::runtime( + "live streaming invocation was cancelled before execution", + ))); + return CommandOutcome::Continue; + } + let idempotency_key = invocation + .idempotency_key() + .cloned() + .unwrap_or_else(IdempotencyKey::fresh); + let mut invocation_context = invocation.invocation_context(); + let result = async { + self.store + .data_mut() + .set_current_idempotency_key(idempotency_key.clone()) + .await; + Self::extend_invocation_context( + &mut invocation_context, + &idempotency_key, + &invocation, + &self.owned_agent_id.agent_id(), + &self.parent.parsed_agent_id, + ); + self.store + .data_mut() + .set_current_invocation_context(invocation_context) + .await?; + let lowered = lower_invocation( + invocation, + &self.store.data().component_metadata().metadata, + self.parent.parsed_agent_id.as_ref(), + )?; + invoke_live_streaming_rpc(lowered, self.store, self.instance, sender, cancellation) + .await + } + .await; + if let Err(error) = result { + tracing::warn!("live streaming invocation failed: {error}"); + } + live_streaming_invocation_outcome(self.parent.agent_mode()) + } + /// Process an external queued worker invocation - this is either an exported function invocation /// or a manual update request (which involves invoking the exported save-snapshot functions, so /// it is a special case of the exported function invocation). @@ -1771,6 +1852,16 @@ fn failed_agent_invocation_outcome( } } +fn live_streaming_invocation_outcome(agent_mode: AgentMode) -> CommandOutcome { + if agent_mode == AgentMode::Ephemeral { + CommandOutcome::BreakInnerLoopAndArchiveEphemeralOplog(RetryDecision::None) + } else { + // Live streams are not replayable yet. Recreate the Store from the last durable boundary + // before any later command can observe unjournaled guest-memory changes. + CommandOutcome::BreakInnerLoop(RetryDecision::Immediate) + } +} + fn should_cleanup_terminal_ephemeral_invocation( agent_mode: AgentMode, is_agent_component: bool, @@ -1829,8 +1920,8 @@ fn snapshot_action_at( mod tests { use super::{ CommandOutcome, PeriodicSnapshotAction, failed_agent_invocation_outcome, - periodic_snapshot_failure_outcome, snapshot_action_at, snapshot_baseline_timestamp, - successful_agent_invocation_outcome, + live_streaming_invocation_outcome, periodic_snapshot_failure_outcome, snapshot_action_at, + snapshot_baseline_timestamp, successful_agent_invocation_outcome, }; use crate::worker::RetryDecision; use crate::worker::invocation::InvokeResult; @@ -1960,4 +2051,20 @@ mod tests { CommandOutcome::BreakInnerLoop(RetryDecision::None) ); } + + #[test] + fn durable_live_streaming_invocation_always_reconstructs_the_store() { + assert_eq!( + live_streaming_invocation_outcome(AgentMode::Durable), + CommandOutcome::BreakInnerLoop(RetryDecision::Immediate) + ); + } + + #[test] + fn ephemeral_live_streaming_invocation_archives_its_ephemeral_oplog() { + assert_eq!( + live_streaming_invocation_outcome(AgentMode::Ephemeral), + CommandOutcome::BreakInnerLoopAndArchiveEphemeralOplog(RetryDecision::None) + ); + } } diff --git a/golem-worker-executor/src/worker/mod.rs b/golem-worker-executor/src/worker/mod.rs index 59f484158d..9e31a84a29 100644 --- a/golem-worker-executor/src/worker/mod.rs +++ b/golem-worker-executor/src/worker/mod.rs @@ -896,6 +896,7 @@ impl Worker { FinalWorkerState::Unloaded { startup_failure: None, }, + PendingLiveInvocationDisposition::Fail, ) .await; @@ -930,8 +931,13 @@ impl Worker { self.status_flusher.begin_delete().await; self.status_checkpointer.begin_delete().await; let error = WorkerExecutorError::invalid_request("Worker is being deleted"); - self.stop_internal(false, Some(error), FinalWorkerState::Deleting) - .await; + self.stop_internal( + false, + Some(error), + FinalWorkerState::Deleting, + PendingLiveInvocationDisposition::Fail, + ) + .await; Ok(()) } @@ -1498,7 +1504,16 @@ impl Worker { invocation: AgentInvocation, idempotency_key: IdempotencyKey, ) -> Result { - match self.clone().invoke(invocation).await? { + let result = self.clone().invoke(invocation).await?; + self.await_invocation_result(idempotency_key, result).await + } + + pub(crate) async fn await_invocation_result( + self: Arc, + idempotency_key: IdempotencyKey, + result: ResultOrSubscription, + ) -> Result { + match result { ResultOrSubscription::Finished(Ok(output)) => Ok(output), ResultOrSubscription::Finished(Err(err)) => Err(err), ResultOrSubscription::Pending(subscription) => { @@ -2033,6 +2048,7 @@ impl Worker { FinalWorkerState::Unloaded { startup_failure: None, }, + PendingLiveInvocationDisposition::Fail, ) .await; drop(instance_guard); @@ -2179,11 +2195,38 @@ impl Worker { } } - /// Acquire storage semaphore permits for a write operation. - /// Called from `DurableWorkerCtx::acquire_filesystem_space` in live mode only. + /// Acquire temporary storage semaphore permits for an unpersisted execution. + /// Called from the durable host in live mode only. /// Returns `NodeOutOfFilesystemStorage` if the executor pool is exhausted. /// /// Should only be called from the invocation loop. + pub async fn acquire_unpersisted_filesystem_storage_space( + &self, + new_bytes: u64, + ) -> anyhow::Result> { + if new_bytes == 0 { + return Ok(None); + } + match &*self.instance.lock().await { + WorkerInstance::Running(_) => { + if let Some(permit) = self + .active_workers() + .try_acquire_filesystem_storage(new_bytes) + .await + { + self.desired_extra_filesystem_storage + .store(0, Ordering::Relaxed); + Ok(Some(permit)) + } else { + self.desired_extra_filesystem_storage + .store(new_bytes, Ordering::Relaxed); + Err(anyhow!(GolemSpecificWasmTrap::NodeOutOfFilesystemStorage)) + } + } + _ => Ok(None), + } + } + pub async fn acquire_filesystem_storage_space(&self, new_bytes: u64) -> anyhow::Result<()> { match &mut *self.instance.lock().await { WorkerInstance::Running(running) => { @@ -2595,6 +2638,59 @@ impl Worker { receiver.await.unwrap() } + /// Runs an awaited agent invocation through the dedicated live-streaming + /// path. Unlike normal invocations, this path is intentionally not added to + /// the durable queue and never serializes its result into the oplog. + pub async fn invoke_live_streaming( + self: Arc, + invocation: AgentInvocation, + cancellation: tokio_util::sync::CancellationToken, + ) -> Result { + self.enqueue_live_streaming(invocation, cancellation) + .await? + .result() + .await + } + + pub async fn enqueue_live_streaming( + self: Arc, + invocation: AgentInvocation, + cancellation: tokio_util::sync::CancellationToken, + ) -> Result { + let instance_guard = self.lock_non_stopping_worker().await; + if instance_guard.is_deleting() { + return Err(WorkerExecutorError::invalid_request( + "Cannot invoke a deleting worker", + )); + } + if let Some(err) = instance_guard.startup_failure() { + return Err(err.clone()); + } + + let (sender, receiver) = oneshot::channel(); + self.queue + .write() + .await + .push_back(QueuedWorkerInvocation::LiveStreamingInvocation { + invocation: Box::new(invocation), + sender, + cancellation: cancellation.clone(), + }); + if let WorkerInstance::Running(running) = &*instance_guard { + running.sender.send(WorkerCommand::WorkAvailable).unwrap(); + } + drop(instance_guard); + tokio::spawn(async move { + if let Err(error) = Worker::start_if_needed(self).await { + tracing::debug!(%error, "Failed to start worker for live streaming invocation"); + } + }); + Ok(LiveStreamingInvocationHandle { + receiver, + cancellation: Some(cancellation), + }) + } + /// Appends an oplog entry without forcing a durable commit. Callers that /// require ordering must await the append before exposing subsequent work. pub async fn add_to_oplog(&self, entry: OplogEntry) -> OplogIndex { @@ -2925,6 +3021,7 @@ impl Worker { called_from_invocation_loop: bool, fail_pending_invocations: Option, final_state: FinalWorkerState, + pending_live_invocations: PendingLiveInvocationDisposition, ) { let mut instance_guard = self.instance.lock().await; @@ -2934,6 +3031,7 @@ impl Worker { called_from_invocation_loop, fail_pending_invocations, final_state, + pending_live_invocations, ) .await; @@ -2950,6 +3048,7 @@ impl Worker { // Only respected when this is the call that triggered the stop fail_pending_invocations: Option, final_state: FinalWorkerState, + pending_live_invocations: PendingLiveInvocationDisposition, ) -> StopResult { // Temporarily set the instance to unloaded so we can work with the old value. // This is not visible to anyone as long as we are holding the lock. @@ -2966,6 +3065,13 @@ impl Worker { self.fail_pending_invocations(error.clone()).await; } **instance_guard = final_state.into_instance(); + if let WorkerInstance::Unloaded { startup_failure } = &**instance_guard { + self.resolve_pending_queue_on_unload( + startup_failure.as_ref(), + pending_live_invocations, + ) + .await; + } StopResult::Stopped } WorkerInstance::WaitingForPermit(_) => { @@ -2975,8 +3081,11 @@ impl Worker { crate::metrics::workers::dec_worker_waiting_for_memory(); **instance_guard = final_state.into_instance(); if let WorkerInstance::Unloaded { startup_failure } = &**instance_guard { - self.resolve_pending_readiness_awaiters_on_stop(startup_failure.as_ref()) - .await; + self.resolve_pending_queue_on_unload( + startup_failure.as_ref(), + pending_live_invocations, + ) + .await; } StopResult::Stopped } @@ -2990,6 +3099,9 @@ impl Worker { StopResult::Stopped } WorkerInstance::Stopping(mut stopping) => { + if pending_live_invocations == PendingLiveInvocationDisposition::Fail { + stopping.pending_live_invocations = PendingLiveInvocationDisposition::Fail; + } // If we're stopping for deletion, upgrade the final state if matches!(final_state, FinalWorkerState::Deleting) { stopping.final_state = FinalWorkerState::Deleting; @@ -3039,8 +3151,11 @@ impl Worker { self.release_linear_memory_grant(); **instance_guard = final_state.into_instance(); if let WorkerInstance::Unloaded { startup_failure } = &**instance_guard { - self.resolve_pending_readiness_awaiters_on_stop(startup_failure.as_ref()) - .await; + self.resolve_pending_queue_on_unload( + startup_failure.as_ref(), + pending_live_invocations, + ) + .await; } StopResult::Stopped } else { @@ -3053,6 +3168,7 @@ impl Worker { **instance_guard = WorkerInstance::Stopping(StoppingWorker { notify: notify.clone(), final_state, + pending_live_invocations, }); StopResult::NeedsWaitForLoopExit { run_loop_handle, @@ -3094,20 +3210,25 @@ impl Worker { instance_guard = self.instance.lock().await; } - match std::mem::replace( + let pending_live_invocations = match std::mem::replace( &mut *instance_guard, WorkerInstance::Unloaded { startup_failure: None, }, ) { WorkerInstance::Stopping(stopping) => { + let pending_live_invocations = stopping.pending_live_invocations; *instance_guard = stopping.final_state.into_instance(); + pending_live_invocations } other => panic!("expected Stopping, got {other:?}"), - } + }; if let WorkerInstance::Unloaded { startup_failure } = &*instance_guard { - self.resolve_pending_readiness_awaiters_on_stop(startup_failure.as_ref()) - .await; + self.resolve_pending_queue_on_unload( + startup_failure.as_ref(), + pending_live_invocations, + ) + .await; } drop(instance_guard); @@ -3121,6 +3242,19 @@ impl Worker { /// worker suspends itself mid-invocation, as debugging workers do as soon as their replay /// goes live. Waiters observe the startup failure if there is one, otherwise a successful /// stop. All other queued items are kept for the next start. + async fn resolve_pending_queue_on_unload( + &self, + startup_failure: Option<&WorkerExecutorError>, + pending_live_invocations: PendingLiveInvocationDisposition, + ) { + self.resolve_pending_readiness_awaiters_on_stop(startup_failure) + .await; + if pending_live_invocations == PendingLiveInvocationDisposition::Fail { + self.fail_pending_live_streaming_invocations_on_stop(startup_failure) + .await; + } + } + async fn resolve_pending_readiness_awaiters_on_stop( &self, startup_failure: Option<&WorkerExecutorError>, @@ -3140,6 +3274,19 @@ impl Worker { } } + async fn fail_pending_live_streaming_invocations_on_stop( + &self, + startup_failure: Option<&WorkerExecutorError>, + ) { + let error = startup_failure.cloned().unwrap_or_else(|| { + WorkerExecutorError::runtime( + "worker stopped before processing the live streaming invocation", + ) + }); + let mut queue = self.queue.write().await; + fail_live_streaming_queue_entries(&mut queue, &error); + } + async fn fail_pending_invocations(&self, error: WorkerExecutorError) { let queued_items = self.queue.write().await.drain(..).collect::>(); let mut origins = self.external_invocation_origins.write().await; @@ -3159,6 +3306,9 @@ impl Worker { QueuedWorkerInvocation::AwaitReadyToProcessCommands { sender } => { let _ = sender.send(Err(error.clone())); } + QueuedWorkerInvocation::LiveStreamingInvocation { sender, .. } => { + let _ = sender.send(Err(error.clone())); + } QueuedWorkerInvocation::SaveSnapshot => {} } } @@ -3218,6 +3368,7 @@ impl Worker { FinalWorkerState::Unloaded { startup_failure: None, }, + PendingLiveInvocationDisposition::Fail, ) .await; let instance_guard = self.instance.lock().await; @@ -3240,6 +3391,7 @@ impl Worker { FinalWorkerState::Unloaded { startup_failure: None, }, + PendingLiveInvocationDisposition::Preserve, ) .await; if let Some(delay) = delay { @@ -4590,10 +4742,17 @@ impl FinalWorkerState { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PendingLiveInvocationDisposition { + Fail, + Preserve, +} + #[derive(Debug)] struct StoppingWorker { notify: OneShotEvent, final_state: FinalWorkerState, + pending_live_invocations: PendingLiveInvocationDisposition, } #[derive(Debug, Clone)] @@ -5045,6 +5204,57 @@ mod tests { assert!(terminal(InterruptKind::Interrupt(Timestamp::now_utc()))); assert!(terminal(InterruptKind::Suspend(Timestamp::now_utc()))); } + + #[test] + async fn terminal_stop_fails_queued_live_invocations_without_dropping_other_work() { + let (sender, receiver) = oneshot::channel(); + let mut queue = VecDeque::from([ + QueuedWorkerInvocation::LiveStreamingInvocation { + invocation: Box::new(AgentInvocation::AgentInitialization { + idempotency_key: IdempotencyKey::fresh(), + input: golem_common::schema::SchemaValue::Record { fields: Vec::new() }, + invocation_context: InvocationContextStack::fresh(), + principal: Principal::anonymous(), + }), + sender, + cancellation: tokio_util::sync::CancellationToken::new(), + }, + QueuedWorkerInvocation::SaveSnapshot, + ]); + let error = WorkerExecutorError::runtime("worker generation stopped"); + + fail_live_streaming_queue_entries(&mut queue, &error); + + assert!(matches!( + receiver.await, + Ok(Err(WorkerExecutorError::Runtime { details, .. })) + if details == "worker generation stopped" + )); + assert!(matches!( + queue.pop_front(), + Some(QueuedWorkerInvocation::SaveSnapshot) + )); + assert!(queue.is_empty()); + } + + #[test] + async fn completed_live_result_is_not_overridden_by_later_worker_state() { + let (sender, receiver) = oneshot::channel(); + let cancellation = tokio_util::sync::CancellationToken::new(); + sender + .send(Ok(golem_common::schema::SchemaValue::U64(42))) + .unwrap(); + let handle = LiveStreamingInvocationHandle { + receiver, + cancellation: Some(cancellation.clone()), + }; + + assert_eq!( + handle.result().await.unwrap(), + golem_common::schema::SchemaValue::U64(42) + ); + assert!(!cancellation.is_cancelled()); + } } #[derive(Clone, Debug, Eq, PartialEq, Hash)] @@ -5192,9 +5402,56 @@ pub enum QueuedWorkerInvocation { AwaitReadyToProcessCommands { sender: oneshot::Sender>, }, + LiveStreamingInvocation { + invocation: Box, + sender: oneshot::Sender>, + cancellation: tokio_util::sync::CancellationToken, + }, SaveSnapshot, } +fn fail_live_streaming_queue_entries( + queue: &mut VecDeque, + error: &WorkerExecutorError, +) { + let items = queue.drain(..).collect::>(); + for item in items { + match item { + QueuedWorkerInvocation::LiveStreamingInvocation { sender, .. } => { + let _ = sender.send(Err(error.clone())); + } + other => queue.push_back(other), + } + } +} + +pub struct LiveStreamingInvocationHandle { + receiver: oneshot::Receiver>, + cancellation: Option, +} + +impl LiveStreamingInvocationHandle { + pub async fn result( + mut self, + ) -> Result { + let result = (&mut self.receiver).await.map_err(|_| { + WorkerExecutorError::runtime( + "live streaming invocation ended before publishing a result", + ) + })?; + self.cancellation = None; + result + } +} + +impl Drop for LiveStreamingInvocationHandle { + fn drop(&mut self) { + if let Some(cancellation) = self.cancellation.take() { + cancellation.cancel(); + } + } +} + #[allow(clippy::large_enum_variant)] pub enum ResultOrSubscription { Finished(Result), diff --git a/golem-worker-executor/tests/agent.rs b/golem-worker-executor/tests/agent.rs index a1e48beb01..3f5a1f41be 100644 --- a/golem-worker-executor/tests/agent.rs +++ b/golem-worker-executor/tests/agent.rs @@ -14,19 +14,24 @@ use crate::Tracing; -use golem_api_grpc::proto::golem::workerexecutor; +use golem_api_grpc::proto::golem::worker::InvocationStart; use golem_common::model::agent::AgentMode; +use golem_common::model::component::ComponentRevision; use golem_common::model::oplog::{OplogIndex, PublicOplogEntry}; use golem_common::model::worker::AgentConfigEntryDto; -use golem_common::model::{AgentId, IdempotencyKey, InvocationStatus}; +use golem_common::model::{AgentId, IdempotencyKey}; use golem_common::schema::SchemaValue; use golem_common::{agent_id, data_value}; use golem_test_framework::dsl::TestDsl; +use golem_worker_executor::storage::scheduler::SchedulerStorage; +use golem_worker_executor::storage::scheduler::sqlite::SqliteSchedulerStorage; use golem_worker_executor_test_utils::{ - LastUniqueId, PrecompiledComponent, TestContext, WorkerExecutorTestDependencies, start, + LastUniqueId, PrecompiledComponent, TestContext, WorkerExecutorTestDependencies, + scheduler_sqlite_storage_config, start, }; use pretty_assertions::assert_eq; -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::path::{Path, PathBuf}; use test_r::{inherit_test_dep, test, timeout}; inherit_test_dep!(WorkerExecutorTestDependencies); @@ -35,6 +40,10 @@ inherit_test_dep!( #[tagged_as("agent_rpc")] PrecompiledComponent ); +inherit_test_dep!( + #[tagged_as("agent_rpc_rust")] + PrecompiledComponent +); inherit_test_dep!( #[tagged_as("constructor_parameter_echo_unnamed")] PrecompiledComponent @@ -98,6 +107,209 @@ async fn agent_self_rpc_is_not_allowed( Ok(()) } +#[test] +#[timeout("60s")] +#[tracing::instrument] +async fn streaming_schedule_is_rejected_without_creating_or_queueing_a_worker( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let scheduler = + SqliteSchedulerStorage::configured(&scheduler_sqlite_storage_config(deps, &context)) + .await + .map_err(anyhow::Error::msg)?; + let agent_id = agent_id!("StreamingRpcTarget", "rejected-schedule"); + let worker_id = AgentId::from_agent_id(component.id, &agent_id).map_err(anyhow::Error::msg)?; + let (_, input) = data_value!(vec![1_u32, 2, 3]).into_parts(); + let input: golem_api_grpc::proto::golem::schema::SchemaValue = + input.try_into().map_err(anyhow::Error::msg)?; + let component_id = component.id.to_string(); + let blobs_before = files_below(&deps.blob_storage_root())? + .into_iter() + .filter(|path| path.to_string_lossy().contains(&component_id)) + .collect::>(); + + for schedule_at in [ + None, + Some(prost_types::Timestamp { + seconds: chrono::Utc::now().timestamp() + 1, + nanos: 0, + }), + ] { + let error = executor + .invoke_agent_session(InvocationStart { + agent_id: Some(worker_id.clone().into()), + method_name: Some("produce".to_string()), + input: Some(input.clone()), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32, + schedule_at, + idempotency_key: Some(IdempotencyKey::fresh().into()), + component_owner_account_id: Some(component.account_id.into()), + environment_id: Some(component.environment_id.into()), + auth_ctx: Some(executor.auth_ctx().into()), + context: None, + principal: None, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + config: Vec::new(), + }) + .await + .expect_err("scheduled streaming invocation must be rejected"); + assert!( + error + .to_string() + .contains("require an attached Await invocation session"), + "unexpected error: {error}" + ); + } + + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + assert_eq!(executor.get_worker_metadata_opt(&worker_id).await?, None); + let assignment = golem_common::model::ShardAssignment { + number_of_shards: 1, + shard_ids: HashSet::from([golem_common::model::ShardId::new(0)]), + }; + assert_eq!( + scheduler + .count_due(chrono::Utc::now() + chrono::Duration::days(1), &assignment) + .await?, + 0, + "rejection must not create an immediate or delayed scheduled action" + ); + assert_eq!( + files_below(&deps.blob_storage_root())? + .into_iter() + .filter(|path| path.to_string_lossy().contains(&component_id)) + .collect::>(), + blobs_before, + "rejection must not upload an invocation or result payload" + ); + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn invocation_classification_uses_the_existing_workers_component_revision( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc")] agent_rpc: &PrecompiledComponent, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + + let streaming_component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let streaming_agent_id = agent_id!("StreamingRpcTarget", "pinned-streaming-revision"); + let streaming_worker_id = executor + .start_agent(&streaming_component.id, streaming_agent_id.clone()) + .await?; + executor + .update_component(&streaming_component.id, &agent_rpc.wasm_name) + .await?; + + let (_, input) = data_value!(vec![1_u32, 2, 3]).into_parts(); + let error = executor + .invoke_agent_session(InvocationStart { + agent_id: Some(streaming_worker_id.clone().into()), + method_name: Some("produce".to_string()), + input: Some(input.try_into().map_err(anyhow::Error::msg)?), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32, + schedule_at: None, + idempotency_key: Some(IdempotencyKey::fresh().into()), + component_owner_account_id: Some(streaming_component.account_id.into()), + environment_id: Some(streaming_component.environment_id.into()), + auth_ctx: Some(executor.auth_ctx().into()), + context: None, + principal: None, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + config: Vec::new(), + }) + .await + .expect_err("the old streaming schema must still reject scheduling"); + assert!( + error + .to_string() + .contains("require an attached Await invocation session"), + "unexpected error: {error}" + ); + let streaming_metadata = executor.get_worker_metadata(&streaming_worker_id).await?; + assert_eq!( + streaming_metadata.component_revision, + ComponentRevision::INITIAL + ); + assert_eq!(streaming_metadata.pending_invocation_count, 0); + + let stream_free_component = executor + .component_dep(&context.default_environment_id, agent_rpc) + .store() + .await?; + let stream_free_agent_id = agent_id!("TestAgent", "pinned-stream-free-revision"); + let stream_free_worker_id = executor + .start_agent(&stream_free_component.id, stream_free_agent_id.clone()) + .await?; + executor + .update_component(&stream_free_component.id, &agent_rpc_rust.wasm_name) + .await?; + + let output = executor + .invoke_and_await_agent( + &stream_free_component, + &stream_free_agent_id, + "run", + data_value!(0_f64), + ) + .await?; + assert_eq!( + output.into_return_value(), + Some(SchemaValue::List { + elements: Vec::new() + }) + ); + let stream_free_metadata = executor.get_worker_metadata(&stream_free_worker_id).await?; + assert_eq!( + stream_free_metadata.component_revision, + ComponentRevision::INITIAL + ); + Ok(()) +} + +fn files_below(root: &Path) -> anyhow::Result> { + fn visit(root: &Path, current: &Path, files: &mut HashSet) -> anyhow::Result<()> { + if !current.exists() { + return Ok(()); + } + for entry in std::fs::read_dir(current)? { + let path = entry?.path(); + if path.is_dir() { + visit(root, &path, files)?; + } else { + files.insert(path.strip_prefix(root)?.to_path_buf()); + } + } + Ok(()) + } + + let mut files = HashSet::new(); + visit(root, root, &mut files)?; + Ok(files) +} + #[test] #[tracing::instrument] async fn agent_await_parallel_rpc_calls( @@ -388,12 +600,14 @@ async fn immediate_scheduled_ephemeral_invocation_reuses_completed_result( AgentId::from_agent_id(component.id, &final_agent_id).map_err(anyhow::Error::msg)?; executor - .client - .clone() - .invoke_agent(workerexecutor::v1::InvokeAgentRequest { + .invoke_agent_session(InvocationStart { agent_id: Some(worker_id.into()), method_name: Some("changeAndGet".to_string()), - method_parameters: Some(SchemaValue::Tuple { elements: vec![] }.into()), + input: Some( + SchemaValue::Record { fields: vec![] } + .try_into() + .map_err(anyhow::Error::msg)?, + ), mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32, schedule_at: None, idempotency_key: Some(idempotency_key.into()), @@ -402,8 +616,9 @@ async fn immediate_scheduled_ephemeral_invocation_reuses_completed_result( auth_ctx: Some(executor.auth_ctx().into()), context: None, principal: None, - freshness_disposition: workerexecutor::v1::InvocationFreshnessDisposition::MayExist - as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, config: Vec::new(), }) .await?; @@ -437,13 +652,11 @@ async fn ephemeral_invocation_lookup_does_not_create_unknown_agent( let worker_id = AgentId::from_agent_id(component.id, &final_agent_id).map_err(anyhow::Error::msg)?; - let response = executor - .client - .clone() - .invoke_agent(workerexecutor::v1::InvokeAgentRequest { + executor + .invoke_agent_session(InvocationStart { agent_id: Some(worker_id.clone().into()), method_name: None, - method_parameters: None, + input: None, mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup as i32, schedule_at: None, idempotency_key: Some(idempotency_key.into()), @@ -452,24 +665,12 @@ async fn ephemeral_invocation_lookup_does_not_create_unknown_agent( auth_ctx: Some(executor.auth_ctx().into()), context: None, principal: None, - freshness_disposition: workerexecutor::v1::InvocationFreshnessDisposition::MayExist - as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, config: Vec::new(), }) - .await? - .into_inner(); - - let success = match response.result { - Some(workerexecutor::v1::invoke_agent_response::Result::Success(success)) => success, - other => anyhow::bail!("unexpected lookup response: {other:?}"), - }; - assert_eq!( - success.status, - Some( - golem_api_grpc::proto::golem::worker::InvocationStatus::from(InvocationStatus::Unknown,) - as i32 - ) - ); + .await?; assert_eq!(executor.get_worker_metadata_opt(&worker_id).await?, None); Ok(()) @@ -502,12 +703,14 @@ async fn scheduled_ephemeral_invocation_uses_schedule_time_component_revision( AgentId::from_agent_id(component.id, &final_agent_id).map_err(anyhow::Error::msg)?; executor - .client - .clone() - .invoke_agent(workerexecutor::v1::InvokeAgentRequest { + .invoke_agent_session(InvocationStart { agent_id: Some(worker_id.clone().into()), method_name: Some("changeAndGet".to_string()), - method_parameters: Some(SchemaValue::Tuple { elements: vec![] }.into()), + input: Some( + SchemaValue::Record { fields: vec![] } + .try_into() + .map_err(anyhow::Error::msg)?, + ), mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32, schedule_at: Some(prost_types::Timestamp { seconds: chrono::Utc::now().timestamp() + 3, @@ -519,8 +722,9 @@ async fn scheduled_ephemeral_invocation_uses_schedule_time_component_revision( auth_ctx: Some(executor.auth_ctx().into()), context: None, principal: None, - freshness_disposition: workerexecutor::v1::InvocationFreshnessDisposition::MayExist - as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, config: Vec::new(), }) .await?; diff --git a/golem-worker-executor/tests/api.rs b/golem-worker-executor/tests/api.rs index 6d8964fbf1..a30650102b 100644 --- a/golem-worker-executor/tests/api.rs +++ b/golem-worker-executor/tests/api.rs @@ -263,7 +263,7 @@ async fn delete_interrupts_long_rpc_call( .invoke_agent( &component, &agent_id, - "long-rpc-call", + "longRpcCall", data_value!(600000f64), // 10 minutes ) .await?; @@ -1931,7 +1931,7 @@ async fn trying_to_use_a_wasm_that_wasmtime_cannot_load_provides_good_error_mess // trying to invoke the previously created worker let result = executor - .invoke_and_await_agent(&component, &agent_id, "run", data_value!()) + .invoke_and_await_agent(&component, &agent_id, "sleep_for", data_value!(0.0f64)) .await; let err = result.expect_err("Expected ComponentParseFailed error"); diff --git a/golem-worker-executor/tests/rpc.rs b/golem-worker-executor/tests/rpc.rs index f8dda85c50..ff47763c93 100644 --- a/golem-worker-executor/tests/rpc.rs +++ b/golem-worker-executor/tests/rpc.rs @@ -13,18 +13,41 @@ // limitations under the License. use crate::Tracing; -use golem_common::model::oplog::OplogIndex; -use golem_common::model::{AgentStatus, IdempotencyKey, PromiseId}; +use async_trait::async_trait; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::schema::{RecordValue, SchemaValueStreamReference, schema_value}; +use golem_api_grpc::proto::golem::worker::{ + InvocationFailureKind, InvocationRejectionReason, InvocationRequest, InvocationResponse, + InvocationStart, ResumeAttach, StreamCancel, StreamCancelReason, StreamCancelRole, + invocation_request, invocation_response, invocation_session_completion, + invocation_session_result, +}; +use golem_common::model::account::AccountId; +use golem_common::model::agent::ParsedAgentId; +use golem_common::model::card::{AgentResourcePattern, AgentVerb}; +use golem_common::model::component::ComponentDto; +use golem_common::model::oplog::{OplogIndex, PublicAgentInvocation, PublicOplogEntry}; +use golem_common::model::{AgentId, AgentStatus, IdempotencyKey, OwnedAgentId, PromiseId}; use golem_common::schema::schema_value::ResultValuePayload; -use golem_common::schema::{FromSchema, SchemaValue}; +use golem_common::schema::{FromSchema, SchemaValue, TypedSchemaValue}; use golem_common::{agent_id, data_value}; +use golem_service_base::model::auth::AuthCtx; use golem_test_framework::dsl::TestDsl; +use golem_worker_executor::services::direct_invocation_auth::{ + DirectInvocationAuthService, EnvironmentOwnerAccountId, +}; +use golem_worker_executor::services::rpc::RpcError; +use golem_worker_executor::worker::EvictionClass; use golem_worker_executor_test_utils::{ - LastUniqueId, PrecompiledComponent, TestContext, WorkerExecutorTestDependencies, start, + LastUniqueId, PrecompiledComponent, TestContext, TestExecutorOverrides, TestWorkerExecutor, + WorkerExecutorTestDependencies, start, start_with_overrides, }; use pretty_assertions::assert_eq; +use std::sync::Arc; use std::time::Duration; -use test_r::{inherit_test_dep, test}; +use test_r::{inherit_test_dep, test, timeout}; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; use tracing::Instrument; inherit_test_dep!(WorkerExecutorTestDependencies); @@ -41,8 +64,1046 @@ inherit_test_dep!( #[tagged_as("agent_counters")] PrecompiledComponent ); +inherit_test_dep!( + #[tagged_as("large_dynamic_memory")] + PrecompiledComponent +); inherit_test_dep!(Tracing); +struct DenyDirectInvocationAuth; + +#[async_trait] +impl DirectInvocationAuthService for DenyDirectInvocationAuth { + async fn check( + &self, + _caller_account_id: AccountId, + _owned_agent_id: &OwnedAgentId, + _verb: AgentVerb, + _resource: AgentResourcePattern, + _auth_ctx: &AuthCtx, + ) -> Result { + Err(RpcError::Denied { + details: "direct invocation denied before schema lookup".to_string(), + }) + } +} + +#[test] +#[timeout("60s")] +#[tracing::instrument] +async fn resume_attach_is_terminally_rejected_without_finish( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let idempotency_key = Some(IdempotencyKey::fresh().into()); + let (requests, receiver) = mpsc::channel(1); + requests + .send(InvocationRequest { + request: Some(invocation_request::Request::ResumeAttach(ResumeAttach { + idempotency_key: idempotency_key.clone(), + })), + }) + .await?; + + let mut responses = executor + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let response = responses + .message() + .await? + .ok_or_else(|| anyhow::anyhow!("resume-attach returned no rejection"))?; + assert!(matches!( + response.response, + Some(invocation_response::Response::Rejected(rejected)) + if rejected.reason == InvocationRejectionReason::ResumeUnsupported as i32 + && rejected.idempotency_key == idempotency_key + )); + assert!(responses.message().await?.is_none()); + Ok(()) +} + +#[test] +#[timeout("60s")] +#[tracing::instrument] +async fn invalid_start_is_rejected_before_acceptance( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let (requests, receiver) = mpsc::channel(1); + requests + .send(InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(golem_api_grpc::proto::golem::worker::AgentId { + component_id: None, + name: "agent".to_string(), + }), + method_name: Some("run".to_string()), + input: Some(golem_api_grpc::proto::golem::schema::SchemaValue { + value: Some( + golem_api_grpc::proto::golem::schema::schema_value::Value::U8Value(1), + ), + }), + idempotency_key: Some(IdempotencyKey::fresh().into()), + auth_ctx: None, + ..Default::default() + })), + }) + .await?; + + let mut responses = executor + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let response = responses + .message() + .await? + .ok_or_else(|| anyhow::anyhow!("invalid invocation start returned no rejection"))?; + + match response.response { + Some(invocation_response::Response::Rejected(_)) => {} + other => { + panic!("invalid invocation start must be rejected before acceptance, got {other:?}") + } + } + assert!(responses.message().await?.is_none()); + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn output_consumer_cancel_after_result_remains_a_valid_terminal_session( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let agent_id = agent_id!("StreamingRpcTarget", "cancel-sibling-output"); + let worker_agent_id = AgentId::from_agent_id(component.id, &agent_id) + .map_err(|error| anyhow::anyhow!("invalid agent id: {error}"))?; + let (_, input) = data_value!().into_parts(); + let key = Some(IdempotencyKey::fresh().into()); + let start = InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(worker_agent_id.into()), + method_name: Some("produce_siblings".to_string()), + input: Some(input.try_into().map_err(anyhow::Error::msg)?), + idempotency_key: key, + auth_ctx: Some(executor.auth_ctx().into()), + environment_id: Some(component.environment_id.into()), + component_owner_account_id: Some(component.account_id.into()), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + ..Default::default() + })), + }; + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&start) + .map_err(anyhow::Error::msg)?; + let (requests, receiver) = mpsc::channel(8); + requests.send(start).await?; + let mut responses = executor + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let mut cancellation_sent = false; + + while let Some(response) = responses.message().await? { + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + match response.response { + Some(invocation_response::Response::Result(result)) => { + let value = match result.result { + Some(invocation_session_result::Result::MethodResult(value)) => value, + other => anyhow::bail!("expected a method result, got {other:?}"), + }; + let stream_id = match value.value { + Some(schema_value::Value::TupleValue(tuple)) => match tuple.elements.first() { + Some(golem_api_grpc::proto::golem::schema::SchemaValue { + value: Some(schema_value::Value::StreamReference(reference)), + }) => reference.stream_id, + other => anyhow::bail!("expected first sibling stream, got {other:?}"), + }, + other => anyhow::bail!("expected sibling tuple result, got {other:?}"), + }; + let cancel = InvocationRequest { + request: Some(invocation_request::Request::StreamCancel(StreamCancel { + stream_id, + offset: 0, + role: StreamCancelRole::OutputConsumer as i32, + reason: StreamCancelReason::Cancelled as i32, + details: Some("consumer stopped reading".to_string()), + })), + }; + state + .validate_trusted_request(&cancel) + .map_err(anyhow::Error::msg)?; + requests.send(cancel).await?; + cancellation_sent = true; + } + Some(invocation_response::Response::Finished(finished)) => { + assert!(cancellation_sent, "session finished before cancellation"); + if let Some(invocation_session_completion::Outcome::Failure(failure)) = + finished.outcome + { + assert_ne!( + failure.kind, + InvocationFailureKind::Protocol as i32, + "a validator-approved output-consumer cancellation is not a protocol error: {}", + failure.message + ); + } + assert!(state.is_complete()); + assert!(responses.message().await?.is_none()); + return Ok(()); + } + Some(invocation_response::Response::Rejected(rejected)) => { + anyhow::bail!("invocation rejected: {}", rejected.error) + } + _ => {} + } + } + + anyhow::bail!("invocation response closed without InvocationFinished") +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn malformed_request_after_streaming_result_terminalizes_open_streams( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let agent_id = agent_id!("StreamingRpcTarget", "malformed-after-result"); + let worker_agent_id = AgentId::from_agent_id(component.id, &agent_id) + .map_err(|error| anyhow::anyhow!("invalid agent id: {error}"))?; + let input = golem_api_grpc::proto::golem::schema::SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![golem_api_grpc::proto::golem::schema::SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 1 }, + )), + }], + })), + }; + let start = InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(worker_agent_id.clone().into()), + method_name: Some("transform".to_string()), + input: Some(input), + idempotency_key: Some(IdempotencyKey::fresh().into()), + auth_ctx: Some(executor.auth_ctx().into()), + environment_id: Some(component.environment_id.into()), + component_owner_account_id: Some(component.account_id.into()), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + ..Default::default() + })), + }; + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&start) + .map_err(anyhow::Error::msg)?; + let (requests, receiver) = mpsc::channel(8); + requests.send(start.clone()).await?; + let mut responses = executor + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let mut expected_stream_ids = Vec::new(); + let mut terminal_stream_ids = Vec::new(); + let mut accepted = false; + let mut result_received = false; + + while let Some(response) = responses.message().await? { + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + match response.response { + Some(invocation_response::Response::Accepted(_)) => { + assert!(!accepted, "invocation was accepted more than once"); + accepted = true; + } + Some(invocation_response::Response::Result(result)) => { + assert!(accepted, "streaming result preceded volatile acceptance"); + result_received = true; + let value = match result.result { + Some(invocation_session_result::Result::MethodResult(value)) => value, + other => anyhow::bail!("expected a method result, got {other:?}"), + }; + expected_stream_ids = match value.value { + Some(schema_value::Value::StreamReference(reference)) => { + vec![reference.stream_id] + } + other => anyhow::bail!("expected transform stream, got {other:?}"), + }; + requests.send(start.clone()).await?; + } + Some(invocation_response::Response::OutputEnd(end)) => { + terminal_stream_ids.push(end.stream_id); + } + Some(invocation_response::Response::OutputError(error)) => { + terminal_stream_ids.push(error.stream_id); + } + Some(invocation_response::Response::Finished(finished)) => { + assert!(accepted, "post-enqueue failure preceded acceptance"); + assert!( + result_received, + "protocol failure preceded the streaming result" + ); + let failure = match finished.outcome { + Some(invocation_session_completion::Outcome::Failure(failure)) => failure, + other => anyhow::bail!("expected protocol failure, got {other:?}"), + }; + assert_eq!(failure.kind, InvocationFailureKind::Protocol as i32); + expected_stream_ids.sort_unstable(); + terminal_stream_ids.sort_unstable(); + assert_eq!(terminal_stream_ids, expected_stream_ids); + assert!(state.is_complete()); + assert!(responses.message().await?.is_none()); + let oplog = executor + .get_oplog(&worker_agent_id, OplogIndex::INITIAL) + .await?; + assert!( + !oplog.iter().any(|entry| matches!( + &entry.entry, + PublicOplogEntry::AgentInvocationStarted(started) + if matches!( + &started.invocation, + PublicAgentInvocation::AgentMethodInvocation(method) + if method.method_name == "transform" + ) + )), + "volatile streaming invocation leaked into the oplog" + ); + return Ok(()); + } + Some(invocation_response::Response::Rejected(rejected)) => { + anyhow::bail!("invocation rejected: {}", rejected.error) + } + _ => {} + } + } + + anyhow::bail!("invocation response closed without InvocationFinished") +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn disconnect_immediately_after_acceptance_cancels_volatile_invocation( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let agent_id = agent_id!("StreamingRpcTarget", "disconnect-after-acceptance"); + let worker_agent_id = executor + .start_agent(&component.id, agent_id.clone()) + .await?; + let input = golem_api_grpc::proto::golem::schema::SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![golem_api_grpc::proto::golem::schema::SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 1 }, + )), + }], + })), + }; + let start = InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(worker_agent_id.clone().into()), + method_name: Some("transform".to_string()), + input: Some(input), + idempotency_key: Some(IdempotencyKey::fresh().into()), + auth_ctx: Some(executor.auth_ctx().into()), + environment_id: Some(component.environment_id.into()), + component_owner_account_id: Some(component.account_id.into()), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + ..Default::default() + })), + }; + let (requests, receiver) = mpsc::channel(8); + requests.send(start).await?; + let mut responses = executor + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let first = responses + .message() + .await? + .ok_or_else(|| anyhow::anyhow!("streaming invocation ended before acceptance"))?; + assert!(matches!( + first.response, + Some(invocation_response::Response::Accepted(_)) + )); + + drop(requests); + drop(responses); + + let ping = tokio::time::timeout( + Duration::from_secs(30), + invoke_agent_session(&executor, &component, &agent_id, "ping", data_value!()), + ) + .await + .map_err(|_| anyhow::anyhow!("cancelled live invocation blocked the next invocation"))?? + .map_err(anyhow::Error::msg)?; + assert_eq!(ping, SchemaValue::U64(42)); + + let oplog = executor + .get_oplog(&worker_agent_id, OplogIndex::INITIAL) + .await?; + assert!( + !oplog.iter().any(|entry| matches!( + &entry.entry, + PublicOplogEntry::AgentInvocationStarted(started) + if matches!( + &started.invocation, + PublicAgentInvocation::AgentMethodInvocation(method) + if method.method_name == "transform" + ) + )), + "disconnected live invocation leaked into the oplog" + ); + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn reacquire_permits_restart_preserves_accepted_queued_live_invocation( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("large_dynamic_memory")] large_dynamic_memory: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + const EXECUTOR_MEMORY_BYTES: u64 = 32 * 1024 * 1024; + const GROWTH_MIB: u64 = 30; + const QUEUE_GATE_MILLIS: u64 = 5_000; + + let context = TestContext::new(last_unique_id); + let executor = start_with_overrides( + deps, + &context, + TestExecutorOverrides { + configure: Some(Arc::new(|config| { + config.memory.system_memory_override = Some(EXECUTOR_MEMORY_BYTES); + config.memory.worker_memory_ratio = 1.0; + config.memory.component_size_coefficient = 0.0; + })), + ..Default::default() + }, + ) + .await?; + let component = executor + .component_dep(&context.default_environment_id, large_dynamic_memory) + .store() + .await?; + + let victim_agent = agent_id!("LargeDynamicMemoryAgent", "live-queue-reacquire-victim"); + let victim_worker = executor + .start_agent(&component.id, victim_agent.clone()) + .await?; + let victim_owned = OwnedAgentId::new(context.default_environment_id, &victim_worker); + tokio::time::timeout(Duration::from_secs(10), async { + while executor.worker_eviction_class(&victim_owned).await != Some(EvictionClass::LoadedIdle) + { + tokio::task::yield_now().await; + } + }) + .await + .expect("the co-resident worker must become idle and evictable"); + + let target_agent = agent_id!("LargeDynamicMemoryAgent", "live-queue-reacquire-target"); + let target_worker = executor + .start_agent(&component.id, target_agent.clone()) + .await?; + let target_owned = OwnedAgentId::new(context.default_environment_id, &target_worker); + let victim_bytes = executor.worker_memory_requirement(&victim_owned).await?; + let target_bytes = executor.worker_memory_requirement(&target_owned).await?; + let growth_bytes = GROWTH_MIB * 1024 * 1024; + assert!(victim_bytes + target_bytes <= EXECUTOR_MEMORY_BYTES); + assert!(target_bytes + growth_bytes <= EXECUTOR_MEMORY_BYTES); + assert!(victim_bytes + target_bytes + growth_bytes > EXECUTOR_MEMORY_BYTES); + + let blocker_executor = executor.clone(); + let blocker_component = component.clone(); + let blocker_agent = target_agent.clone(); + let blocker = tokio::spawn(async move { + blocker_executor + .invoke_and_await_agent( + &blocker_component, + &blocker_agent, + "run_with_memory_and_work", + data_value!(0u64, QUEUE_GATE_MILLIS), + ) + .await + }); + executor + .wait_for_status( + &target_worker, + AgentStatus::Running, + Duration::from_secs(10), + ) + .await?; + + executor + .invoke_agent( + &component, + &target_agent, + "run_with_memory_and_work", + data_value!(GROWTH_MIB, 0u64), + ) + .await?; + + let (mut state, _frames, mut inbound) = open_invocation_session( + &executor, + &component, + &target_agent, + "run_with_memory_and_work", + data_value!(0u64, 0u64), + ) + .await?; + let accepted = tokio::time::timeout(Duration::from_secs(3), inbound.message()) + .await + .map_err(|_| anyhow::anyhow!("live invocation was not accepted while durable work ran"))?? + .ok_or_else(|| anyhow::anyhow!("live invocation ended before acceptance"))?; + state + .validate_response(&accepted) + .map_err(anyhow::Error::msg)?; + assert!(matches!( + accepted.response, + Some(invocation_response::Response::Accepted(_)) + )); + assert!( + !blocker.is_finished(), + "the queue gate ended before the live invocation was accepted" + ); + + let blocker_result = blocker.await??; + assert_eq!(blocker_result.into_typed::()?, 0); + + let live_result = tokio::time::timeout( + Duration::from_secs(60), + receive_invocation_session(&mut state, &mut inbound), + ) + .await + .map_err(|_| { + anyhow::anyhow!("accepted live invocation was stranded by permit reacquisition restart") + })?? + .map_err(anyhow::Error::msg)?; + assert_eq!(live_result, SchemaValue::U64(0)); + assert!( + !executor.worker_is_loaded(&victim_owned).await, + "the durable growth must force permit reacquisition and evict the idle worker" + ); + + Ok(()) +} + +async fn invoke_agent_session( + executor: &TestWorkerExecutor, + component: &ComponentDto, + agent_id: &ParsedAgentId, + method_name: &str, + params: TypedSchemaValue, +) -> anyhow::Result> { + let (mut state, _frames, mut inbound) = + open_invocation_session(executor, component, agent_id, method_name, params).await?; + receive_invocation_session(&mut state, &mut inbound).await +} + +async fn open_invocation_session( + executor: &TestWorkerExecutor, + component: &ComponentDto, + agent_id: &ParsedAgentId, + method_name: &str, + params: TypedSchemaValue, +) -> anyhow::Result<( + InvocationSessionState, + mpsc::Sender, + tonic::Streaming, +)> { + let worker_agent_id = AgentId::from_agent_id(component.id, agent_id) + .map_err(|error| anyhow::anyhow!("invalid agent id: {error}"))?; + let (_, input) = params.into_parts(); + let input = input.try_into().map_err(anyhow::Error::msg)?; + let (frames, receiver) = mpsc::channel(8); + let request = InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(worker_agent_id.into()), + method_name: Some(method_name.to_string()), + input: Some(input), + idempotency_key: Some(IdempotencyKey::fresh().into()), + context: None, + auth_ctx: Some(executor.auth_ctx().into()), + principal: None, + environment_id: Some(component.environment_id.into()), + config: Vec::new(), + component_owner_account_id: Some(component.account_id.into()), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + schedule_at: None, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + })), + }; + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&request) + .map_err(anyhow::Error::msg)?; + frames.send(request).await?; + let inbound = executor + .client + .clone() + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + Ok((state, frames, inbound)) +} + +async fn receive_invocation_session( + state: &mut InvocationSessionState, + inbound: &mut tonic::Streaming, +) -> anyhow::Result> { + let mut result = None; + let mut terminal = None; + while let Some(response) = inbound.message().await? { + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + match response.response { + Some(invocation_response::Response::Accepted(_)) => {} + Some(invocation_response::Response::Rejected(rejected)) => { + terminal = Some(Err(rejected.error)); + } + Some(invocation_response::Response::Result(value)) => { + if result.is_some() { + anyhow::bail!("invocation session returned more than one result"); + } + result = match value.result { + Some(invocation_session_result::Result::MethodResult(value)) => { + Some(value.try_into().map_err(anyhow::Error::msg)?) + } + Some(invocation_session_result::Result::NoResult(_)) | None => { + anyhow::bail!("invocation session returned no method result") + } + }; + } + Some(invocation_response::Response::Finished(finished)) => { + terminal = Some(match finished.outcome { + Some(invocation_session_completion::Outcome::Success(_)) => { + result.take().map(Ok).ok_or_else(|| { + anyhow::anyhow!("invocation session ended without a result") + })? + } + Some(invocation_session_completion::Outcome::Failure(failure)) => { + Err(failure.message) + } + None => anyhow::bail!("invocation session completion has no outcome"), + }); + } + Some(other) => { + anyhow::bail!("unexpected outer invocation session frame: {other:?}") + } + None => anyhow::bail!("empty outer invocation session frame"), + } + } + terminal.ok_or_else(|| anyhow::anyhow!("invocation session response ended before completion")) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn generated_rust_client_streaming_rpc_e2e( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let caller_agent_id = agent_id!("StreamingRpcCaller", "generated_streaming_rpc_e2e"); + let caller = executor + .start_agent(&component.id, caller_agent_id.clone()) + .await?; + + let result = invoke_agent_session( + &executor, + &component, + &caller_agent_id, + "run", + data_value!(), + ) + .await? + .map_err(anyhow::Error::msg)?; + let SchemaValue::Record { fields } = result else { + panic!("expected streaming RPC report record"); + }; + assert_eq!(fields.len(), 10); + assert_eq!( + fields[0], + SchemaValue::List { + elements: vec![ + SchemaValue::U32(1), + SchemaValue::U32(2), + SchemaValue::U32(3) + ] + } + ); + assert_eq!( + fields[1], + SchemaValue::List { + elements: vec![ + SchemaValue::U32(4), + SchemaValue::U32(5), + SchemaValue::U32(6) + ] + } + ); + assert_eq!( + fields[2], + SchemaValue::List { + elements: vec![ + SchemaValue::U32(70), + SchemaValue::U32(80), + SchemaValue::U32(90) + ] + } + ); + assert_eq!( + fields[3], + SchemaValue::List { + elements: vec![ + SchemaValue::String("left".to_string()), + SchemaValue::String("right".to_string()) + ] + } + ); + assert_eq!( + fields[4], + SchemaValue::List { + elements: vec![SchemaValue::U32(10), SchemaValue::U32(11)] + } + ); + assert_eq!( + fields[5], + SchemaValue::List { + elements: vec![ + SchemaValue::String("first".to_string()), + SchemaValue::String("second".to_string()) + ] + } + ); + assert_eq!( + fields[6], + SchemaValue::List { + elements: vec![ + SchemaValue::List { + elements: vec![SchemaValue::U32(1), SchemaValue::U32(2)] + }, + SchemaValue::List { + elements: vec![ + SchemaValue::U32(3), + SchemaValue::U32(4), + SchemaValue::U32(5) + ] + } + ] + } + ); + assert_eq!( + fields[7], + SchemaValue::List { + elements: vec![ + SchemaValue::String("a".to_string()), + SchemaValue::String("b".to_string()) + ] + } + ); + assert_eq!( + fields[8], + SchemaValue::List { + elements: (0..64).map(SchemaValue::U32).collect() + } + ); + assert_eq!(fields[9], SchemaValue::U64(42)); + + let producer_error = invoke_agent_session( + &executor, + &component, + &caller_agent_id, + "call_producer_error", + data_value!(), + ) + .await? + .expect_err("producer stream error must fail the invocation session"); + assert!( + producer_error.contains("Component trapped"), + "unexpected producer error: {producer_error}" + ); + + let stream_free_caller_id = agent_id!("StreamingRpcCaller", "stream_free_after_stream_error"); + executor + .start_agent(&component.id, stream_free_caller_id.clone()) + .await?; + let first = executor + .invoke_and_await_agent( + &component, + &stream_free_caller_id, + "call_stream_free", + data_value!(), + ) + .await? + .into_typed::()?; + let second = executor + .invoke_and_await_agent( + &component, + &stream_free_caller_id, + "call_stream_free", + data_value!(), + ) + .await? + .into_typed::()?; + assert_eq!((first, second), (1, 2)); + executor.check_oplog_is_queryable(&caller).await?; + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn output_drop_cancels_target_blocked_on_stream_input( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let caller_agent_id = agent_id!("StreamingRpcCaller", "blocked-stream-input-cancellation"); + executor + .start_agent(&component.id, caller_agent_id.clone()) + .await?; + + let result = invoke_agent_session( + &executor, + &component, + &caller_agent_id, + "cancel_transform_blocked_on_input", + data_value!(), + ) + .await? + .map_err(anyhow::Error::msg)?; + assert_eq!( + result, + SchemaValue::U64(42), + "dropping output must cancel a target invocation blocked on its input before ping can run" + ); + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn durable_agent_live_await_streaming_is_allowed( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let caller_agent_id = agent_id!("StreamingRpcCaller", "durable-live-await"); + + let result = invoke_agent_session( + &executor, + &component, + &caller_agent_id, + "run", + data_value!(), + ) + .await? + .map_err(anyhow::Error::msg)?; + + let SchemaValue::Record { fields } = result else { + panic!("expected streaming RPC report record"); + }; + assert_eq!(fields.len(), 10); + assert_eq!(fields[9], SchemaValue::U64(42)); + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn direct_rpc_classification_uses_the_existing_target_revision( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc")] agent_rpc: &PrecompiledComponent, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start(deps, &context).await?; + + let ts_component = executor + .component_dep(&context.default_environment_id, agent_rpc) + .store() + .await?; + let ts_caller_id = agent_id!("TestAgent", "pinned-ts-rpc-caller"); + let ts_target_id = agent_id!("ChildAgent", 0_f64); + executor + .start_agent(&ts_component.id, ts_caller_id.clone()) + .await?; + executor.start_agent(&ts_component.id, ts_target_id).await?; + executor + .update_component(&ts_component.id, &agent_rpc_rust.wasm_name) + .await?; + + let ts_result = executor + .invoke_and_await_agent(&ts_component, &ts_caller_id, "run", data_value!(1_f64)) + .await?; + assert_eq!( + ts_result.into_return_value(), + Some(SchemaValue::List { + elements: vec![SchemaValue::F64(0.0)] + }) + ); + + let rust_component = executor + .component_dep(&context.default_environment_id, agent_rpc_rust) + .store() + .await?; + let rust_name = "pinned-rust-rpc-target"; + let rust_caller_id = agent_id!("StreamingRpcCaller", rust_name); + let rust_target_id = agent_id!("StreamingRpcTarget", rust_name); + executor + .start_agent(&rust_component.id, rust_caller_id.clone()) + .await?; + executor + .start_agent(&rust_component.id, rust_target_id) + .await?; + executor + .update_component(&rust_component.id, &agent_rpc.wasm_name) + .await?; + + let rust_result = executor + .invoke_and_await_agent( + &rust_component, + &rust_caller_id, + "call_stream_free", + data_value!(), + ) + .await? + .into_typed::()?; + assert_eq!(rust_result, 1); + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn direct_rpc_authorizes_before_execution_revision_lookup( + last_unique_id: &LastUniqueId, + deps: &WorkerExecutorTestDependencies, + #[tagged_as("agent_rpc")] agent_rpc: &PrecompiledComponent, + #[tagged_as("agent_rpc_rust")] agent_rpc_rust: &PrecompiledComponent, + _tracing: &Tracing, +) -> anyhow::Result<()> { + let context = TestContext::new(last_unique_id); + let executor = start_with_overrides( + deps, + &context, + TestExecutorOverrides { + create_direct_invocation_auth: Some(Arc::new(|| Arc::new(DenyDirectInvocationAuth))), + ..Default::default() + }, + ) + .await?; + let component = executor + .component_dep(&context.default_environment_id, agent_rpc) + .store() + .await?; + let caller_id = agent_id!("TestAgent", "denied-rpc-caller"); + executor + .start_agent(&component.id, caller_id.clone()) + .await?; + executor + .update_component(&component.id, &agent_rpc_rust.wasm_name) + .await?; + + let error = executor + .invoke_and_await_agent(&component, &caller_id, "run", data_value!(1_f64)) + .await + .expect_err("direct RPC must be denied"); + assert!( + error + .to_string() + .contains("direct invocation denied before schema lookup"), + "unexpected error: {error}" + ); + Ok(()) +} + #[test] #[tracing::instrument] async fn rust_rpc_with_payload( diff --git a/golem-worker-service/src/api/agents.rs b/golem-worker-service/src/api/agents.rs index 5ba653c76b..009397afe0 100644 --- a/golem-worker-service/src/api/agents.rs +++ b/golem-worker-service/src/api/agents.rs @@ -1,7 +1,16 @@ use crate::api::common::ApiEndpointError; -use crate::service::auth::AuthService; -use crate::service::worker::WorkerService; +use crate::service::auth::{AuthService, AuthServiceError}; +use crate::service::worker::{ + WorkerService, WorkerServiceError, validate_public_session_schema_value, +}; use chrono::{DateTime, Utc}; +use futures::{SinkExt, StreamExt}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::worker::{ + InvocationRejected, InvocationRejectionReason, InvocationRequest, InvocationResponse, + PublicInvocationRequest, PublicInvocationStart, input_stream_item, invocation_request, + invocation_response, public_invocation_request, +}; use golem_common::base_model::api; use golem_common::model::agent::AgentTypeName; use golem_common::model::application::ApplicationName; @@ -9,13 +18,18 @@ use golem_common::model::component::ComponentRevision; use golem_common::model::environment::EnvironmentName; use golem_common::model::worker::AgentConfigEntryDto; use golem_common::model::{AgentId, IdempotencyKey}; -use golem_common::recorded_http_api_request; use golem_common::schema::{SchemaValue, TypedSchemaValue}; +use golem_common::{SafeDisplay, recorded_http_api_request}; use golem_service_base::api_tags::ApiTags; -use golem_service_base::model::auth::GolemSecurityScheme; +use golem_service_base::clients::registry::RegistryServiceError; +use golem_service_base::model::auth::{AuthCtx, GolemSecurityScheme}; +use poem::web::websocket::{ + BoxWebSocketUpgraded, CloseCode, Message, WebSocket, WebSocketConfig, WebSocketStream, +}; use poem_openapi::param::Header; use poem_openapi::payload::Json; use poem_openapi_derive::{Enum, Object, OpenApi}; +use prost::Message as ProstMessage; use serde::{Deserialize, Serialize}; use std::sync::Arc; use tracing::Instrument; @@ -23,6 +37,13 @@ use uuid::Uuid; type Result = std::result::Result; +const INVOCATION_SESSION_CHANNEL_CAPACITY: usize = 16; +const INVOCATION_SESSION_MAX_MESSAGE_SIZE: usize = 32 * 1024 * 1024; +const INVOCATION_SESSION_STAGING_ADMISSION_TIMEOUT: std::time::Duration = + std::time::Duration::from_secs(30); +const INVOCATION_SESSION_WRITE_PROGRESS_TIMEOUT: std::time::Duration = + std::time::Duration::from_secs(30); + pub struct AgentsApi { worker_service: Arc, auth_service: Arc, @@ -80,6 +101,31 @@ impl AgentsApi { record.result(response).map(Json) } + /// Invoke an agent through an attached live streaming session + #[oai( + path = "/invoke-agent-session", + method = "get", + operation_id = "invoke_agent_session" + )] + async fn invoke_agent_session( + &self, + websocket: WebSocket, + token: GolemSecurityScheme, + ) -> Result { + let auth = self.auth_service.authenticate_token(token.secret()).await?; + let worker_service = self.worker_service.clone(); + + Ok(websocket + .config(invocation_session_websocket_config()) + .on_upgrade(Box::new(move |socket| { + Box::pin(proxy_public_invocation_session( + socket, + worker_service, + auth, + )) + }))) + } + #[oai(path = "/create-agent", method = "post", operation_id = "create_agent")] async fn create_agent( &self, @@ -106,6 +152,566 @@ impl AgentsApi { } } +fn invocation_session_websocket_config() -> WebSocketConfig { + WebSocketConfig::default() + .max_message_size(Some(INVOCATION_SESSION_MAX_MESSAGE_SIZE)) + .max_frame_size(Some(INVOCATION_SESSION_MAX_MESSAGE_SIZE)) + .max_write_buffer_size(2 * INVOCATION_SESSION_MAX_MESSAGE_SIZE) +} + +#[derive(Debug)] +enum PublicSessionMessage { + Request(Box), + Ping(Vec), + Pong, + Close, +} + +enum InitialPublicInvocation { + Start( + Box, + Option, + ), + Closed, + WriterStopped, +} + +#[derive(Debug)] +struct PublicSessionFrameError { + close_code: CloseCode, + reason: String, +} + +fn decode_public_session_message( + message: Message, +) -> std::result::Result { + match message { + Message::Binary(bytes) => PublicInvocationRequest::decode(bytes.as_slice()) + .map(Box::new) + .map(PublicSessionMessage::Request) + .map_err(|error| PublicSessionFrameError { + close_code: CloseCode::Protocol, + reason: format!("malformed invocation request: {error}"), + }), + Message::Text(_) => Err(PublicSessionFrameError { + close_code: CloseCode::Unsupported, + reason: "invocation sessions accept binary protobuf messages only".to_string(), + }), + Message::Ping(payload) => Ok(PublicSessionMessage::Ping(payload)), + Message::Pong(_) => Ok(PublicSessionMessage::Pong), + Message::Close(_) => Ok(PublicSessionMessage::Close), + } +} + +async fn proxy_public_invocation_session( + socket: WebSocketStream, + worker_service: Arc, + auth: AuthCtx, +) { + let (websocket_sink, mut websocket_stream) = socket.split(); + let (websocket_sender, websocket_receiver) = + tokio::sync::mpsc::channel(INVOCATION_SESSION_CHANNEL_CAPACITY); + let mut writer = tokio::spawn(forward_websocket_messages( + websocket_sink, + websocket_receiver, + INVOCATION_SESSION_WRITE_PROGRESS_TIMEOUT, + )); + let mut state = InvocationSessionState::default(); + + let (start, idempotency_key) = match receive_public_invocation_start_while_writing( + &mut websocket_stream, + &websocket_sender, + &mut state, + &mut writer, + ) + .await + { + InitialPublicInvocation::Start(start, idempotency_key) => (*start, idempotency_key), + InitialPublicInvocation::Closed => { + drop(websocket_sender); + let _ = writer.await; + return; + } + InitialPublicInvocation::WriterStopped => return, + }; + + let (request_sender, request_receiver) = + tokio::sync::mpsc::channel(INVOCATION_SESSION_CHANNEL_CAPACITY); + let tail = tokio_stream::wrappers::ReceiverStream::new(request_receiver); + let responses = match worker_service + .invoke_public_agent_session(start, Box::pin(tail), auth) + .await + { + Ok(responses) => responses, + Err(error) => { + let response = rejection_response( + rejection_reason(&error), + error.to_safe_string(), + idempotency_key, + ); + if state.validate_response(&response).is_ok() + && queue_invocation_response(&websocket_sender, response).await + { + queue_websocket_close(&websocket_sender, CloseCode::Normal, "session rejected") + .await; + } + drop(websocket_sender); + let _ = writer.await; + return; + } + }; + + let state = Arc::new(tokio::sync::Mutex::new(state)); + let requests = tokio::spawn(forward_public_requests( + websocket_stream, + request_sender, + websocket_sender.clone(), + state.clone(), + )); + let responses = tokio::spawn(forward_internal_responses( + responses, + websocket_sender.clone(), + state, + )); + drop(websocket_sender); + + supervise_public_invocation_session(requests, responses, writer).await; +} + +async fn receive_public_invocation_start_while_writing( + websocket_stream: &mut S, + websocket_sender: &tokio::sync::mpsc::Sender, + state: &mut InvocationSessionState, + writer: &mut tokio::task::JoinHandle<()>, +) -> InitialPublicInvocation +where + S: futures::Stream> + Unpin, +{ + let receive = receive_public_invocation_start(websocket_stream, websocket_sender, state); + tokio::pin!(receive); + tokio::select! { + result = &mut receive => match result { + Some((start, idempotency_key)) => { + InitialPublicInvocation::Start(Box::new(start), idempotency_key) + } + None => InitialPublicInvocation::Closed, + }, + _ = writer => InitialPublicInvocation::WriterStopped, + } +} + +async fn receive_public_invocation_start( + websocket_stream: &mut S, + websocket_sender: &tokio::sync::mpsc::Sender, + state: &mut InvocationSessionState, +) -> Option<( + PublicInvocationStart, + Option, +)> +where + S: futures::Stream> + Unpin, +{ + let first_request = loop { + match websocket_stream.next().await { + Some(Ok(message)) => match decode_public_session_message(message) { + Ok(PublicSessionMessage::Request(request)) => break *request, + Ok(PublicSessionMessage::Ping(payload)) => { + if matches!( + websocket_sender.try_send(Message::pong(payload)), + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) + ) { + return None; + } + } + Ok(PublicSessionMessage::Pong) => {} + Ok(PublicSessionMessage::Close) => return None, + Err(error) => { + queue_websocket_close(websocket_sender, error.close_code, error.reason).await; + return None; + } + }, + None | Some(Err(_)) => return None, + } + }; + + if let Err(error) = state.validate_public_request(&first_request) { + queue_websocket_close(websocket_sender, CloseCode::Protocol, error).await; + return None; + } + + match first_request.request { + Some(public_invocation_request::Request::Start(start)) => { + let idempotency_key = start.idempotency_key.clone(); + Some((start, idempotency_key)) + } + Some(public_invocation_request::Request::ResumeAttach(resume)) => { + let response = rejection_response( + InvocationRejectionReason::ResumeUnsupported, + "resume-attach is not supported by provisional live sessions".to_string(), + resume.idempotency_key, + ); + if state.validate_response(&response).is_ok() + && queue_invocation_response(websocket_sender, response).await + { + queue_websocket_close(websocket_sender, CloseCode::Normal, "session rejected") + .await; + } + None + } + _ => { + queue_websocket_close( + websocket_sender, + CloseCode::Protocol, + "the first invocation request must be start or resume-attach", + ) + .await; + None + } + } +} + +async fn forward_websocket_messages( + mut websocket_sink: S, + mut websocket_receiver: tokio::sync::mpsc::Receiver, + write_progress_timeout: std::time::Duration, +) where + S: futures::Sink + Unpin, +{ + while let Some(message) = websocket_receiver.recv().await { + let close = matches!(message, Message::Close(_)); + if !matches!( + tokio::time::timeout(write_progress_timeout, websocket_sink.send(message)).await, + Ok(Ok(())) + ) { + return; + } + if close { + return; + } + } +} + +async fn supervise_public_invocation_session( + mut requests: tokio::task::JoinHandle<()>, + mut responses: tokio::task::JoinHandle<()>, + mut writer: tokio::task::JoinHandle<()>, +) { + enum CompletedPump { + Requests, + Responses, + Writer, + } + let completed = tokio::select! { + _ = &mut requests => CompletedPump::Requests, + _ = &mut responses => CompletedPump::Responses, + _ = &mut writer => CompletedPump::Writer, + }; + match completed { + CompletedPump::Requests => { + responses.abort(); + let _ = responses.await; + let _ = writer.await; + } + CompletedPump::Responses => { + requests.abort(); + let _ = requests.await; + let _ = writer.await; + } + CompletedPump::Writer => { + requests.abort(); + responses.abort(); + let _ = requests.await; + let _ = responses.await; + } + } +} + +async fn forward_public_requests( + websocket_stream: S, + request_sender: tokio::sync::mpsc::Sender, + websocket_sender: tokio::sync::mpsc::Sender, + state: Arc>, +) where + S: futures::Stream> + Unpin, +{ + forward_public_requests_with_timeout( + websocket_stream, + request_sender, + websocket_sender, + state, + INVOCATION_SESSION_STAGING_ADMISSION_TIMEOUT, + ) + .await; +} + +async fn forward_public_requests_with_timeout( + websocket_stream: S, + request_sender: tokio::sync::mpsc::Sender, + websocket_sender: tokio::sync::mpsc::Sender, + state: Arc>, + staging_admission_timeout: std::time::Duration, +) where + S: futures::Stream> + Unpin, +{ + let (staging_sender, mut staging_receiver) = + tokio::sync::mpsc::channel(INVOCATION_SESSION_CHANNEL_CAPACITY); + let read_requests = read_public_requests( + websocket_stream, + staging_sender, + websocket_sender.clone(), + state, + staging_admission_timeout, + ); + let forward_requests = async move { + while let Some(request) = staging_receiver.recv().await { + if request_sender.send(request).await.is_err() { + try_queue_websocket_close( + &websocket_sender, + CloseCode::Error, + "internal invocation request stream closed", + ); + return; + } + } + }; + tokio::pin!(read_requests, forward_requests); + tokio::select! { + _ = &mut read_requests => {} + _ = &mut forward_requests => {} + } +} + +async fn read_public_requests( + mut websocket_stream: S, + staging_sender: tokio::sync::mpsc::Sender, + websocket_sender: tokio::sync::mpsc::Sender, + state: Arc>, + staging_admission_timeout: std::time::Duration, +) where + S: futures::Stream> + Unpin, +{ + let mut pending_message = None; + loop { + let message = match pending_message.take() { + Some(message) => message, + None => match websocket_stream.next().await { + Some(message) => message, + None => return, + }, + }; + let message = match message { + Ok(message) => message, + Err(_) => return, + }; + let request = match decode_public_session_message(message) { + Ok(PublicSessionMessage::Request(request)) => *request, + Ok(PublicSessionMessage::Ping(payload)) => { + if matches!( + websocket_sender.try_send(Message::pong(payload)), + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) + ) { + return; + } + continue; + } + Ok(PublicSessionMessage::Pong) => continue, + Ok(PublicSessionMessage::Close) => return, + Err(error) => { + try_queue_websocket_close(&websocket_sender, error.close_code, error.reason); + return; + } + }; + let validation = { + let mut state = state.lock().await; + state.validate_received_public_request(&request) + }; + if let Err(error) = validation { + try_queue_websocket_close(&websocket_sender, CloseCode::Protocol, error); + return; + } + let request = match trusted_tail_request(request) { + Ok(request) => request, + Err(error) => { + try_queue_websocket_close(&websocket_sender, CloseCode::Protocol, error); + return; + } + }; + let admission = + tokio::time::timeout(staging_admission_timeout, staging_sender.send(request)); + tokio::pin!(admission); + tokio::select! { + result = &mut admission => { + if !matches!(result, Ok(Ok(()))) { + return; + } + } + next = websocket_stream.next() => { + match next { + None | Some(Err(_)) | Some(Ok(Message::Close(_))) => return, + Some(message) => pending_message = Some(message), + } + if !matches!(admission.await, Ok(Ok(()))) { + return; + } + } + } + } +} + +async fn forward_internal_responses( + mut responses: S, + websocket_sender: tokio::sync::mpsc::Sender, + state: Arc>, +) where + S: futures::Stream> + Unpin, +{ + while let Some(response) = responses.next().await { + let response = match response { + Ok(response) => response, + Err(_) => break, + }; + let validation = { + let mut state = state.lock().await; + state + .validate_response(&response) + .map(|()| state.is_complete()) + }; + let complete = match validation { + Ok(complete) => complete, + Err(error) => { + queue_websocket_close(&websocket_sender, CloseCode::Protocol, error).await; + return; + } + }; + if websocket_sender + .send(Message::binary(response.encode_to_vec())) + .await + .is_err() + { + return; + } + if complete { + queue_websocket_close(&websocket_sender, CloseCode::Normal, "session complete").await; + return; + } + } + + queue_websocket_close( + &websocket_sender, + CloseCode::Error, + "internal invocation response stream ended before session completion", + ) + .await; +} + +async fn queue_websocket_close( + sender: &tokio::sync::mpsc::Sender, + code: CloseCode, + reason: impl AsRef, +) { + let reason = bounded_close_reason(reason.as_ref()); + let _ = sender.send(Message::close_with(code, reason)).await; +} + +fn try_queue_websocket_close( + sender: &tokio::sync::mpsc::Sender, + code: CloseCode, + reason: impl AsRef, +) { + let reason = bounded_close_reason(reason.as_ref()); + let _ = sender.try_send(Message::close_with(code, reason)); +} + +fn trusted_tail_request( + request: PublicInvocationRequest, +) -> std::result::Result { + let request = match request.request { + Some(public_invocation_request::Request::InputItem(item)) => { + if let Some(input_stream_item::Payload::Value(value)) = &item.payload { + validate_public_session_schema_value(value)?; + } + invocation_request::Request::InputItem(item) + } + Some(public_invocation_request::Request::InputEnd(end)) => { + invocation_request::Request::InputEnd(end) + } + Some(public_invocation_request::Request::StreamCancel(cancel)) => { + invocation_request::Request::StreamCancel(cancel) + } + Some(public_invocation_request::Request::Start(_)) + | Some(public_invocation_request::Request::ResumeAttach(_)) => { + return Err( + "invocation start or resume-attach may only appear as the first request" + .to_string(), + ); + } + None => return Err("invocation request has no payload".to_string()), + }; + Ok(InvocationRequest { + request: Some(request), + }) +} + +fn rejection_reason(error: &WorkerServiceError) -> InvocationRejectionReason { + match error { + WorkerServiceError::TypeChecker(_) + | WorkerServiceError::RegistryServiceError(RegistryServiceError::BadRequest(_)) => { + InvocationRejectionReason::Validation + } + WorkerServiceError::AuthError(AuthServiceError::Unauthorized(_)) + | WorkerServiceError::RegistryServiceError(RegistryServiceError::Unauthorized(_)) + | WorkerServiceError::RegistryServiceError(RegistryServiceError::CouldNotAuthenticate(_)) => { + InvocationRejectionReason::Unauthorized + } + WorkerServiceError::ComponentNotFound(_) + | WorkerServiceError::AgentNotFound(_) + | WorkerServiceError::RegistryServiceError(RegistryServiceError::NotFound(_)) => { + InvocationRejectionReason::NotFound + } + _ => InvocationRejectionReason::Internal, + } +} + +fn rejection_response( + reason: InvocationRejectionReason, + error: String, + idempotency_key: Option, +) -> InvocationResponse { + InvocationResponse { + response: Some(invocation_response::Response::Rejected( + InvocationRejected { + reason: reason as i32, + error, + idempotency_key, + agent_id: None, + component_revision: None, + }, + )), + } +} + +async fn queue_invocation_response( + sender: &tokio::sync::mpsc::Sender, + response: InvocationResponse, +) -> bool { + sender + .send(Message::binary(response.encode_to_vec())) + .await + .is_ok() +} + +fn bounded_close_reason(reason: &str) -> String { + const MAX_CLOSE_REASON_BYTES: usize = 123; + if reason.len() <= MAX_CLOSE_REASON_BYTES { + return reason.to_string(); + } + let mut end = MAX_CLOSE_REASON_BYTES; + while !reason.is_char_boundary(end) { + end -= 1; + } + reason[..end].to_string() +} + #[derive(Debug, Clone, Serialize, Deserialize, Enum)] #[oai(rename_all = "camelCase")] #[serde(rename_all = "camelCase")] @@ -169,15 +775,52 @@ pub struct CreateAgentResponse { #[cfg(test)] mod tests { - use super::{AgentInvocationRequest, CreateAgentRequest}; + use super::{ + AgentInvocationRequest, CreateAgentRequest, INVOCATION_SESSION_MAX_MESSAGE_SIZE, + InitialPublicInvocation, PublicSessionMessage, bounded_close_reason, + decode_public_session_message, forward_internal_responses, forward_public_requests, + forward_public_requests_with_timeout, forward_websocket_messages, + invocation_session_websocket_config, queue_invocation_response, + receive_public_invocation_start, receive_public_invocation_start_while_writing, + rejection_reason, supervise_public_invocation_session, trusted_tail_request, + }; + use crate::service::worker::WorkerServiceError; + use futures::StreamExt; + use golem_api_grpc::invocation_session_protocol::InvocationSessionState; + use golem_api_grpc::proto::golem::worker::{ + AgentId, IdempotencyKey, InputStreamEnd, InputStreamItem, InvocationAccepted, + InvocationRejectionReason, InvocationRequest, InvocationResponse, + InvocationSessionCompletion, InvocationSessionResult, OutputStreamItem, + PublicInvocationRequest, PublicInvocationStart, ResumeAttach, StreamCancel, + StreamCancelReason, StreamCancelRole, input_stream_item, invocation_request, + invocation_response, invocation_session_result, public_invocation_request, + }; + use golem_service_base::clients::registry::RegistryServiceError; + use poem::web::websocket::{CloseCode, Message}; use poem_openapi::types::{ParseFromJSON, ToJSON}; + use prost::Message as ProstMessage; use serde_json::{Value, json}; + use std::sync::Arc; use test_r::test; fn empty_parameter_record() -> Value { json!({ "kind": "record", "value": { "fields": [] } }) } + fn gated_websocket_sink( + permits: tokio::sync::mpsc::Receiver<()>, + delivered: Arc>>, + ) -> impl futures::Sink { + futures::sink::unfold( + (permits, delivered), + |(mut permits, delivered), message| async move { + permits.recv().await.ok_or(())?; + delivered.lock().await.push(message); + Ok((permits, delivered)) + }, + ) + } + #[test] fn create_agent_request_preserves_canonical_schema_value_json() { let parameters = empty_parameter_record(); @@ -232,4 +875,920 @@ mod tests { assert!(CreateAgentRequest::parse_from_json(Some(request_json)).is_err()); } + + #[test] + fn invocation_session_frame_decoder_accepts_one_binary_protobuf_message() { + let request = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputEnd( + InputStreamEnd { + stream_id: 7, + offset: 3, + }, + )), + }; + + let decoded = decode_public_session_message(Message::binary(request.encode_to_vec())) + .expect("binary protobuf frame must decode"); + + assert!(matches!( + decoded, + PublicSessionMessage::Request(request) if matches!(*request, PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputEnd( + InputStreamEnd { + stream_id: 7, + offset: 3, + } + )), + }) + )); + } + + #[test] + fn invocation_session_frame_decoder_rejects_text_and_malformed_protobuf() { + let text = decode_public_session_message(Message::text("not protobuf")) + .expect_err("text frames must be rejected"); + assert_eq!(text.close_code, CloseCode::Unsupported); + + let malformed = decode_public_session_message(Message::binary([0xff, 0xff])) + .expect_err("malformed protobuf must be rejected"); + assert_eq!(malformed.close_code, CloseCode::Protocol); + } + + #[test] + fn invocation_session_websocket_config_bounds_frames_messages_and_writes() { + let config = invocation_session_websocket_config(); + + assert_eq!( + config.max_message_size, + Some(INVOCATION_SESSION_MAX_MESSAGE_SIZE) + ); + assert_eq!( + config.max_frame_size, + Some(INVOCATION_SESSION_MAX_MESSAGE_SIZE) + ); + assert_eq!( + config.max_write_buffer_size, + 2 * INVOCATION_SESSION_MAX_MESSAGE_SIZE + ); + } + + #[test] + fn public_tail_translation_cannot_construct_a_trusted_start() { + let translated = trusted_tail_request(PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputEnd( + InputStreamEnd { + stream_id: 4, + offset: 9, + }, + )), + }) + .unwrap(); + assert!(matches!( + translated.request, + Some(invocation_request::Request::InputEnd(InputStreamEnd { + stream_id: 4, + offset: 9, + })) + )); + + let repeated_start = PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start(Default::default())), + }; + assert!(trusted_tail_request(repeated_start).is_err()); + } + + #[test] + fn public_tail_translation_rejects_recursive_host_capabilities() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue, SecretValue, schema_value, + }; + + let request = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id: 7, + sequence: 0, + payload: Some(input_stream_item::Payload::Value(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![SchemaValue { + value: Some(schema_value::Value::SecretValue( + SecretValue::default(), + )), + }], + })), + })), + }, + )), + }; + + assert!( + trusted_tail_request(request) + .unwrap_err() + .contains("host-managed capability") + ); + } + + #[test] + async fn saturated_websocket_output_does_not_hide_disconnect_after_ping() { + let (websocket_client, websocket_stream) = + futures::channel::mpsc::unbounded::>(); + let (internal_sender, _internal_receiver) = tokio::sync::mpsc::channel(1); + let (websocket_sender, _websocket_receiver) = tokio::sync::mpsc::channel(1); + websocket_sender + .send(Message::Ping(Vec::new())) + .await + .unwrap(); + + let requests = tokio::spawn(forward_public_requests( + websocket_stream, + internal_sender, + websocket_sender, + Arc::new(tokio::sync::Mutex::new(InvocationSessionState::default())), + )); + websocket_client + .unbounded_send(Ok(Message::Ping(Vec::new()))) + .unwrap(); + tokio::task::yield_now().await; + drop(websocket_client); + + tokio::time::timeout(std::time::Duration::from_millis(100), requests) + .await + .expect("saturated WebSocket output hid the client disconnect") + .unwrap(); + } + + #[test] + async fn saturated_input_forwarding_does_not_block_responses_or_tail_cancellation() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue, SchemaValueStreamReference, schema_value, + }; + + let idempotency_key = IdempotencyKey { + value: "input-backpressure".to_string(), + }; + let stream_id = 7; + let start = PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: "app".to_string(), + environment_name: "env".to_string(), + agent_type_name: "agent".to_string(), + constructor_parameters: Some(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue::default())), + }), + method_name: "run".to_string(), + method_parameters: Some(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id }, + )), + }], + })), + }), + idempotency_key: Some(idempotency_key.clone()), + ..Default::default() + }, + )), + }; + let accepted = InvocationResponse { + response: Some(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: Some(AgentId { + component_id: None, + name: "agent".to_string(), + }), + idempotency_key: Some(idempotency_key), + component_revision: Some(1), + }, + )), + }; + let mut state = InvocationSessionState::default(); + state.validate_public_request(&start).unwrap(); + state.validate_response(&accepted).unwrap(); + let state = Arc::new(tokio::sync::Mutex::new(state)); + + let (websocket_client, websocket_stream) = + futures::channel::mpsc::unbounded::>(); + let (internal_sender, mut internal_receiver) = tokio::sync::mpsc::channel(1); + internal_sender + .send(InvocationRequest { request: None }) + .await + .unwrap(); + let (websocket_sender, mut websocket_receiver) = tokio::sync::mpsc::channel(2); + let requests = tokio::spawn(forward_public_requests( + websocket_stream, + internal_sender, + websocket_sender.clone(), + state.clone(), + )); + let input = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id, + sequence: 0, + payload: Some(input_stream_item::Payload::Value( + golem_common::schema::SchemaValue::U32(1) + .try_into() + .unwrap(), + )), + }, + )), + }; + websocket_client + .unbounded_send(Ok(Message::binary(input.encode_to_vec()))) + .unwrap(); + tokio::task::yield_now().await; + + let result = InvocationResponse { + response: Some(invocation_response::Response::Result( + InvocationSessionResult { + result: Some(invocation_session_result::Result::NoResult( + golem_api_grpc::proto::golem::common::Empty {}, + )), + component_revision: Some(1), + agent_id: Some(AgentId { + component_id: None, + name: "agent".to_string(), + }), + idempotency_key: Some(IdempotencyKey { + value: "input-backpressure".to_string(), + }), + ..Default::default() + }, + )), + }; + let responses = tokio::spawn(forward_internal_responses( + tokio_stream::iter([Ok(result.clone())]), + websocket_sender, + state, + )); + let forwarded = + tokio::time::timeout(std::time::Duration::from_secs(1), websocket_receiver.recv()) + .await + .expect("response forwarding blocked behind saturated input") + .expect("response pump closed unexpectedly"); + let Message::Binary(forwarded) = forwarded else { + panic!("response pump did not forward a binary invocation response") + }; + assert_eq!( + InvocationResponse::decode(forwarded.as_slice()).unwrap(), + result + ); + + drop(websocket_client); + tokio::time::timeout(std::time::Duration::from_secs(1), requests) + .await + .expect("disconnect did not terminate the saturated request pump") + .unwrap(); + let _ = responses.await; + assert!(internal_receiver.recv().await.is_some()); + assert!(internal_receiver.recv().await.is_none()); + } + + #[test] + async fn valid_input_burst_applies_backpressure_without_closing_the_session() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue, SchemaValueStreamReference, schema_value, + }; + + let idempotency_key = IdempotencyKey { + value: "lossless-input-backpressure".to_string(), + }; + let stream_id = 7; + let start = PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: "app".to_string(), + environment_name: "env".to_string(), + agent_type_name: "agent".to_string(), + constructor_parameters: Some(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue::default())), + }), + method_name: "run".to_string(), + method_parameters: Some(SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id }, + )), + }), + idempotency_key: Some(idempotency_key.clone()), + ..Default::default() + }, + )), + }; + let accepted = InvocationResponse { + response: Some(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: Some(AgentId { + component_id: None, + name: "agent".to_string(), + }), + idempotency_key: Some(idempotency_key), + component_revision: Some(1), + }, + )), + }; + let mut state = InvocationSessionState::default(); + state.validate_public_request(&start).unwrap(); + state.validate_response(&accepted).unwrap(); + + let (websocket_client, websocket_stream) = + futures::channel::mpsc::unbounded::>(); + let (internal_sender, mut internal_receiver) = tokio::sync::mpsc::channel(1); + internal_sender + .send(InvocationRequest { request: None }) + .await + .unwrap(); + let (websocket_sender, mut websocket_receiver) = tokio::sync::mpsc::channel(1); + let requests = tokio::spawn(forward_public_requests( + websocket_stream, + internal_sender, + websocket_sender, + Arc::new(tokio::sync::Mutex::new(state)), + )); + + for sequence in 0..18 { + let input = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id, + sequence, + payload: Some(input_stream_item::Payload::Value(SchemaValue { + value: Some(schema_value::Value::U8Value(sequence as u32)), + })), + }, + )), + }; + websocket_client + .unbounded_send(Ok(Message::binary(input.encode_to_vec()))) + .unwrap(); + } + + let outbound = tokio::time::timeout( + std::time::Duration::from_millis(100), + websocket_receiver.recv(), + ) + .await; + assert!( + outbound.is_err(), + "a valid input burst must be backpressured, not close the live session; got {outbound:?}" + ); + assert!(!requests.is_finished()); + + assert!(internal_receiver.recv().await.is_some()); + for expected_sequence in 0..18 { + let request = tokio::time::timeout( + std::time::Duration::from_secs(1), + internal_receiver.recv(), + ) + .await + .expect("lossless input forwarding remained blocked after downstream capacity returned") + .expect("internal invocation request stream closed during a valid input burst"); + assert!(matches!( + request.request, + Some(invocation_request::Request::InputItem(InputStreamItem { + stream_id: actual_stream_id, + sequence: actual_sequence, + .. + })) if actual_stream_id == stream_id && actual_sequence == expected_sequence + )); + } + drop(websocket_client); + tokio::time::timeout(std::time::Duration::from_secs(1), requests) + .await + .expect("disconnect did not terminate the drained request pump") + .unwrap(); + assert!(internal_receiver.recv().await.is_none()); + } + + #[test] + async fn saturated_staging_does_not_hide_websocket_disconnect() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue, SchemaValueStreamReference, schema_value, + }; + + let idempotency_key = IdempotencyKey { + value: "saturated-staging-disconnect".to_string(), + }; + let stream_id = 7; + let start = PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: "app".to_string(), + environment_name: "env".to_string(), + agent_type_name: "agent".to_string(), + constructor_parameters: Some(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue::default())), + }), + method_name: "run".to_string(), + method_parameters: Some(SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id }, + )), + }), + idempotency_key: Some(idempotency_key.clone()), + ..Default::default() + }, + )), + }; + let accepted = InvocationResponse { + response: Some(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: Some(AgentId { + component_id: None, + name: "agent".to_string(), + }), + idempotency_key: Some(idempotency_key), + component_revision: Some(1), + }, + )), + }; + let mut state = InvocationSessionState::default(); + state.validate_public_request(&start).unwrap(); + state.validate_response(&accepted).unwrap(); + + let (websocket_client, websocket_stream) = + futures::channel::mpsc::unbounded::>(); + let (internal_sender, mut internal_receiver) = tokio::sync::mpsc::channel(1); + internal_sender + .send(InvocationRequest { request: None }) + .await + .unwrap(); + let (websocket_sender, _websocket_receiver) = tokio::sync::mpsc::channel(1); + let requests = tokio::spawn(forward_public_requests_with_timeout( + websocket_stream, + internal_sender, + websocket_sender, + Arc::new(tokio::sync::Mutex::new(state)), + std::time::Duration::from_millis(50), + )); + + for sequence in 0..19 { + let input = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id, + sequence, + payload: Some(input_stream_item::Payload::Value(SchemaValue { + value: Some(schema_value::Value::U8Value(sequence as u32)), + })), + }, + )), + }; + websocket_client + .unbounded_send(Ok(Message::binary(input.encode_to_vec()))) + .unwrap(); + } + + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + assert!(!requests.is_finished()); + drop(websocket_client); + + tokio::time::timeout(std::time::Duration::from_millis(200), requests) + .await + .expect("saturated staging exceeded its bounded disconnect-detection deadline") + .unwrap(); + assert!(internal_receiver.recv().await.is_some()); + assert!(internal_receiver.recv().await.is_none()); + } + + #[test] + async fn saturated_output_forwarding_does_not_block_output_cancellation() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue, SchemaValueStreamReference, schema_value, + }; + + let idempotency_key = IdempotencyKey { + value: "output-backpressure".to_string(), + }; + let start = PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: "app".to_string(), + environment_name: "env".to_string(), + agent_type_name: "agent".to_string(), + constructor_parameters: Some(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue::default())), + }), + method_name: "run".to_string(), + method_parameters: Some(SchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue::default())), + }), + idempotency_key: Some(idempotency_key.clone()), + ..Default::default() + }, + )), + }; + let agent_id = AgentId { + component_id: None, + name: "agent".to_string(), + }; + let accepted = InvocationResponse { + response: Some(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: Some(agent_id.clone()), + idempotency_key: Some(idempotency_key.clone()), + component_revision: Some(1), + }, + )), + }; + let mut initial_state = InvocationSessionState::default(); + initial_state.validate_public_request(&start).unwrap(); + initial_state.validate_response(&accepted).unwrap(); + + let output_stream_id = 2; + let result = InvocationResponse { + response: Some(invocation_response::Response::Result( + InvocationSessionResult { + result: Some(invocation_session_result::Result::MethodResult( + SchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { + stream_id: output_stream_id, + }, + )), + }, + )), + component_revision: Some(1), + agent_id: Some(agent_id), + idempotency_key: Some(idempotency_key), + ..Default::default() + }, + )), + }; + initial_state.validate_response(&result).unwrap(); + let state = Arc::new(tokio::sync::Mutex::new(initial_state)); + + let output_item = InvocationResponse { + response: Some(invocation_response::Response::OutputItem( + OutputStreamItem { + stream_id: output_stream_id, + offset: 0, + value: Some(SchemaValue { + value: Some(schema_value::Value::U8Value(42)), + }), + }, + )), + }; + + let (websocket_sender, mut websocket_receiver) = tokio::sync::mpsc::channel(1); + websocket_sender + .send(Message::Ping(Vec::new())) + .await + .unwrap(); + let responses = tokio::spawn(forward_internal_responses( + tokio_stream::iter([Ok(output_item)]), + websocket_sender.clone(), + state.clone(), + )); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + assert!( + !responses.is_finished(), + "response forwarding did not saturate the bounded WebSocket queue" + ); + assert!( + state.try_lock().is_ok(), + "response forwarding retained protocol state while blocked by output backpressure" + ); + + let (websocket_client, websocket_stream) = + futures::channel::mpsc::unbounded::>(); + let (internal_sender, mut internal_receiver) = tokio::sync::mpsc::channel(1); + let requests = tokio::spawn(forward_public_requests( + websocket_stream, + internal_sender, + websocket_sender, + state, + )); + let cancel = PublicInvocationRequest { + request: Some(public_invocation_request::Request::StreamCancel( + StreamCancel { + stream_id: output_stream_id, + offset: 0, + role: StreamCancelRole::OutputConsumer as i32, + reason: StreamCancelReason::Cancelled as i32, + details: Some("stop output".to_string()), + }, + )), + }; + websocket_client + .unbounded_send(Ok(Message::binary(cancel.encode_to_vec()))) + .unwrap(); + + let forwarded = tokio::time::timeout( + std::time::Duration::from_millis(100), + internal_receiver.recv(), + ) + .await + .expect("output cancellation was blocked behind saturated response forwarding") + .expect("internal invocation request stream closed unexpectedly"); + assert!(matches!( + forwarded.request, + Some(invocation_request::Request::StreamCancel(StreamCancel { + stream_id, + role, + .. + })) if stream_id == output_stream_id && role == StreamCancelRole::OutputConsumer as i32 + )); + + drop(websocket_client); + drop(websocket_receiver.recv().await); + requests.abort(); + responses.abort(); + } + + #[test] + async fn pre_start_rejection_waits_for_temporarily_backpressured_writer() { + let request = PublicInvocationRequest { + request: Some(public_invocation_request::Request::ResumeAttach( + ResumeAttach { + idempotency_key: Some(IdempotencyKey { + value: "temporary-pre-start-stall".to_string(), + }), + }, + )), + }; + let mut websocket_stream = + tokio_stream::iter([Ok(Message::binary(request.encode_to_vec()))]); + let delivered = Arc::new(tokio::sync::Mutex::new(Vec::new())); + let (permit_sender, permit_receiver) = tokio::sync::mpsc::channel(2); + let sink = Box::pin(gated_websocket_sink(permit_receiver, delivered.clone())); + let (websocket_sender, websocket_receiver) = tokio::sync::mpsc::channel(2); + let writer = tokio::spawn(forward_websocket_messages( + sink, + websocket_receiver, + std::time::Duration::from_secs(5), + )); + + assert!( + receive_public_invocation_start( + &mut websocket_stream, + &websocket_sender, + &mut InvocationSessionState::default(), + ) + .await + .is_none() + ); + drop(websocket_sender); + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + assert!( + !writer.is_finished(), + "temporary pre-start backpressure truncated the rejection" + ); + + permit_sender.send(()).await.unwrap(); + permit_sender.send(()).await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), writer) + .await + .expect("pre-start writer did not drain after output progress resumed") + .unwrap(); + + let delivered = delivered.lock().await; + assert_eq!(delivered.len(), 2); + let Message::Binary(rejected) = &delivered[0] else { + panic!("rejection was not delivered first") + }; + assert!(matches!( + InvocationResponse::decode(rejected.as_slice()) + .unwrap() + .response, + Some(invocation_response::Response::Rejected(_)) + )); + assert!(matches!(delivered[1], Message::Close(_))); + } + + #[test] + async fn pre_start_ping_stall_terminates_when_the_writer_stops() { + let messages = (0..18) + .map(|_| Ok(Message::Ping(Vec::new()))) + .collect::>>(); + let mut websocket_stream = tokio_stream::iter(messages).chain(futures::stream::pending()); + let delivered = Arc::new(tokio::sync::Mutex::new(Vec::new())); + let (permit_sender, permit_receiver) = tokio::sync::mpsc::channel(1); + let sink = Box::pin(gated_websocket_sink(permit_receiver, delivered.clone())); + let (websocket_sender, websocket_receiver) = tokio::sync::mpsc::channel(1); + let mut writer = tokio::spawn(forward_websocket_messages( + sink, + websocket_receiver, + std::time::Duration::from_millis(50), + )); + + let result = tokio::time::timeout( + std::time::Duration::from_millis(500), + receive_public_invocation_start_while_writing( + &mut websocket_stream, + &websocket_sender, + &mut InvocationSessionState::default(), + &mut writer, + ), + ) + .await + .expect("pre-start Ping stall outlived the writer progress deadline"); + assert!(matches!(result, InitialPublicInvocation::WriterStopped)); + assert!(delivered.lock().await.is_empty()); + drop(websocket_sender); + drop(permit_sender); + } + + #[test] + async fn pre_start_rejection_terminates_at_writer_progress_deadline() { + let request = PublicInvocationRequest { + request: Some(public_invocation_request::Request::ResumeAttach( + ResumeAttach { + idempotency_key: Some(IdempotencyKey { + value: "permanent-pre-start-stall".to_string(), + }), + }, + )), + }; + let mut websocket_stream = + tokio_stream::iter([Ok(Message::binary(request.encode_to_vec()))]); + let delivered = Arc::new(tokio::sync::Mutex::new(Vec::new())); + let (permit_sender, permit_receiver) = tokio::sync::mpsc::channel(1); + let sink = Box::pin(gated_websocket_sink(permit_receiver, delivered.clone())); + let (websocket_sender, websocket_receiver) = tokio::sync::mpsc::channel(2); + let writer = tokio::spawn(forward_websocket_messages( + sink, + websocket_receiver, + std::time::Duration::from_millis(50), + )); + + assert!( + receive_public_invocation_start( + &mut websocket_stream, + &websocket_sender, + &mut InvocationSessionState::default(), + ) + .await + .is_none() + ); + drop(websocket_sender); + tokio::time::timeout(std::time::Duration::from_millis(500), writer) + .await + .expect("pre-start writer exceeded its progress deadline") + .unwrap(); + assert!(delivered.lock().await.is_empty()); + drop(permit_sender); + } + + #[test] + async fn response_completion_waits_for_temporarily_backpressured_writer() { + let delivered = Arc::new(tokio::sync::Mutex::new(Vec::new())); + let (permit_sender, permit_receiver) = tokio::sync::mpsc::channel(3); + let sink = Box::pin(gated_websocket_sink(permit_receiver, delivered.clone())); + let (websocket_sender, websocket_receiver) = tokio::sync::mpsc::channel(3); + let writer = tokio::spawn(forward_websocket_messages( + sink, + websocket_receiver, + std::time::Duration::from_secs(5), + )); + let requests = tokio::spawn(futures::future::pending()); + let responses = tokio::spawn(async move { + let accepted = InvocationResponse { + response: Some(invocation_response::Response::Accepted(Default::default())), + }; + let finished = InvocationResponse { + response: Some(invocation_response::Response::Finished( + InvocationSessionCompletion::default(), + )), + }; + websocket_sender + .send(Message::binary(accepted.encode_to_vec())) + .await + .unwrap(); + websocket_sender + .send(Message::binary(finished.encode_to_vec())) + .await + .unwrap(); + websocket_sender + .send(Message::close_with(CloseCode::Normal, "session complete")) + .await + .unwrap(); + }); + + let supervisor = tokio::spawn(supervise_public_invocation_session( + requests, responses, writer, + )); + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + assert!( + !supervisor.is_finished(), + "temporary output backpressure truncated queued semantic responses" + ); + + for _ in 0..3 { + permit_sender.send(()).await.unwrap(); + } + tokio::time::timeout(std::time::Duration::from_secs(2), supervisor) + .await + .expect("writer did not drain after output progress resumed") + .unwrap(); + + let delivered = delivered.lock().await; + assert_eq!(delivered.len(), 3); + let Message::Binary(accepted) = &delivered[0] else { + panic!("accepted response was not delivered first") + }; + assert!(matches!( + InvocationResponse::decode(accepted.as_slice()) + .unwrap() + .response, + Some(invocation_response::Response::Accepted(_)) + )); + let Message::Binary(finished) = &delivered[1] else { + panic!("finished response was not delivered second") + }; + assert!(matches!( + InvocationResponse::decode(finished.as_slice()) + .unwrap() + .response, + Some(invocation_response::Response::Finished(_)) + )); + assert!(matches!(delivered[2], Message::Close(_))); + } + + #[test] + async fn permanently_stalled_writer_terminates_at_progress_deadline() { + let delivered = Arc::new(tokio::sync::Mutex::new(Vec::new())); + let (permit_sender, permit_receiver) = tokio::sync::mpsc::channel(1); + let sink = Box::pin(gated_websocket_sink(permit_receiver, delivered.clone())); + let (websocket_sender, websocket_receiver) = tokio::sync::mpsc::channel(2); + let writer = tokio::spawn(forward_websocket_messages( + sink, + websocket_receiver, + std::time::Duration::from_millis(50), + )); + let requests = tokio::spawn(futures::future::pending()); + let responses = tokio::spawn(async move { + websocket_sender + .send(Message::binary( + InvocationResponse { + response: Some(invocation_response::Response::Finished( + InvocationSessionCompletion::default(), + )), + } + .encode_to_vec(), + )) + .await + .unwrap(); + websocket_sender + .send(Message::close_with(CloseCode::Normal, "session complete")) + .await + .unwrap(); + }); + + tokio::time::timeout( + std::time::Duration::from_millis(500), + supervise_public_invocation_session(requests, responses, writer), + ) + .await + .expect("permanently stalled writer exceeded its progress deadline"); + assert!(delivered.lock().await.is_empty()); + drop(permit_sender); + } + + #[test] + async fn invocation_responses_are_forwarded_as_binary_protobuf_messages() { + let response = InvocationResponse { response: None }; + let (sender, mut receiver) = tokio::sync::mpsc::channel(1); + + assert!(queue_invocation_response(&sender, response.clone()).await); + let message = receiver.recv().await.unwrap(); + let Message::Binary(bytes) = message else { + panic!("invocation response must use a binary WebSocket message") + }; + + assert_eq!( + InvocationResponse::decode(bytes.as_slice()).unwrap(), + response + ); + } + + #[test] + fn websocket_close_reasons_are_utf8_safe_and_protocol_sized() { + let reason = "é".repeat(100); + let bounded = bounded_close_reason(&reason); + + assert!(bounded.len() <= 123); + assert!(reason.starts_with(&bounded)); + } + + #[test] + fn public_resolution_errors_map_to_protocol_rejections() { + assert_eq!( + rejection_reason(&WorkerServiceError::TypeChecker("invalid".to_string())), + InvocationRejectionReason::Validation + ); + assert_eq!( + rejection_reason(&WorkerServiceError::RegistryServiceError( + RegistryServiceError::Unauthorized("forbidden".to_string()) + )), + InvocationRejectionReason::Unauthorized + ); + assert_eq!( + rejection_reason(&WorkerServiceError::RegistryServiceError( + RegistryServiceError::NotFound("missing".to_string()) + )), + InvocationRejectionReason::NotFound + ); + } } diff --git a/golem-worker-service/src/custom_api/call_agent/mod.rs b/golem-worker-service/src/custom_api/call_agent/mod.rs index b2c7a07a05..323c9da2ec 100644 --- a/golem-worker-service/src/custom_api/call_agent/mod.rs +++ b/golem-worker-service/src/custom_api/call_agent/mod.rs @@ -106,7 +106,9 @@ impl CallAgentHandler { }; let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - method_params_value.into(); + method_params_value.try_into().map_err(|error| { + anyhow!("method parameters cannot cross the worker boundary: {error}") + })?; let invocation_context = Some(golem_api_grpc::proto::golem::worker::InvocationContext { parent: None, @@ -133,7 +135,6 @@ impl CallAgentHandler { Vec::new(), AuthCtx::System, proto_principal, - Some(resolved_route.route.environment_id), ) .await?; diff --git a/golem-worker-service/src/grpcapi/worker.rs b/golem-worker-service/src/grpcapi/worker.rs index 5740849494..8d2d2b5fbe 100644 --- a/golem-worker-service/src/grpcapi/worker.rs +++ b/golem-worker-service/src/grpcapi/worker.rs @@ -14,7 +14,11 @@ use super::error::WorkerTraceErrorKind; use super::{bad_request_error, validate_protobuf_agent_id}; -use crate::service::worker::{WorkerService, WorkerServiceError}; +use crate::service::worker::{ + InvocationRequestStream, InvocationResponseStream, WorkerService, WorkerServiceError, +}; +use futures::{FutureExt, Stream, StreamExt, stream}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; use golem_api_grpc::proto::golem::common::Empty; use golem_api_grpc::proto::golem::worker::v1::worker_service_server::WorkerService as GrpcWorkerService; use golem_api_grpc::proto::golem::worker::v1::{ @@ -28,6 +32,10 @@ use golem_api_grpc::proto::golem::worker::v1::{ launch_new_worker_response, process_oplog_entries_response, resume_worker_response, revert_worker_response, update_worker_response, }; +use golem_api_grpc::proto::golem::worker::{ + InvocationRejected, InvocationRejectionReason, InvocationRequest, InvocationResponse, + invocation_request, invocation_response, +}; use golem_common::model::agent::InvocationFreshnessDisposition; use golem_common::model::component::{ComponentId, ComponentRevision}; use golem_common::model::oplog::OplogIndex; @@ -41,6 +49,167 @@ use std::sync::Arc; use tonic::{Request, Response, Status}; use tracing::Instrument; +fn service_failure_stream( + reason: InvocationRejectionReason, + error: String, + idempotency_key: Option, + agent_id: Option, +) -> InvocationResponseStream { + Box::pin(stream::once(async move { + Ok(InvocationResponse { + response: Some(invocation_response::Response::Rejected( + InvocationRejected { + reason: reason as i32, + error, + idempotency_key, + agent_id, + component_revision: None, + }, + )), + }) + })) +} + +fn request_identity( + request: &InvocationRequest, +) -> ( + Option, + Option, +) { + match request.request.as_ref() { + Some(invocation_request::Request::Start(start)) => { + (start.idempotency_key.clone(), start.agent_id.clone()) + } + Some(invocation_request::Request::ResumeAttach(resume)) => { + (resume.idempotency_key.clone(), None) + } + _ => (None, None), + } +} + +fn validated_response_stream( + inbound: S, + state: Arc>, + initial_requests_checked: Option>, +) -> InvocationResponseStream +where + S: Stream> + Send + Unpin + 'static, +{ + Box::pin(stream::unfold( + Some((inbound, state, initial_requests_checked)), + |state| async move { + let (mut inbound, response_state, initial_requests_checked) = state?; + if let Some(initial_requests_checked) = initial_requests_checked { + let _ = initial_requests_checked.await; + } + match inbound.next().await { + Some(Ok(response)) => { + let mut state = response_state.lock().await; + match state.validate_response(&response) { + Ok(()) if state.is_complete() => { + drop(state); + match inbound.next().await { + None => Some((Ok(response), None)), + Some(Ok(response_after_terminal)) => { + let details = response_state + .lock() + .await + .validate_response(&response_after_terminal) + .unwrap_err(); + Some((Err(Status::internal(details)), None)) + } + Some(Err(error)) => Some((Err(error), None)), + } + } + Ok(()) => { + drop(state); + Some((Ok(response), Some((inbound, response_state, None)))) + } + Err(details) => Some((Err(Status::internal(details)), None)), + } + } + Some(Err(error)) => Some((Err(error), None)), + None if response_state.lock().await.is_complete() => None, + None => Some(( + Err(Status::unavailable( + "invocation response transport closed before completion", + )), + None, + )), + } + }, + )) +} + +fn validated_request_tail( + inbound: S, + state: Arc>, +) -> (InvocationRequestStream, tokio::sync::oneshot::Receiver<()>) +where + S: Stream> + Send + Unpin + 'static, +{ + let (validated_tx, validated_rx) = tokio::sync::mpsc::channel(32); + let (initial_requests_checked_tx, initial_requests_checked_rx) = + tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let mut inbound = inbound; + let mut initial_requests_checked_tx = Some(initial_requests_checked_tx); + loop { + let Some(next) = inbound.next().now_or_never() else { + let _ = initial_requests_checked_tx.take().unwrap().send(()); + break; + }; + let Some(next) = next else { + let _ = initial_requests_checked_tx.take().unwrap().send(()); + return; + }; + let request = match next { + Ok(request) => request, + Err(_) => { + let _ = initial_requests_checked_tx.take().unwrap().send(()); + return; + } + }; + let invalid = state + .lock() + .await + .validate_received_trusted_request(&request) + .is_err(); + if validated_tx.send(request).await.is_err() { + let _ = initial_requests_checked_tx.take().unwrap().send(()); + return; + } + if invalid { + let _ = initial_requests_checked_tx.take().unwrap().send(()); + return; + } + } + while let Some(next) = inbound.next().await { + let request = match next { + Ok(request) => request, + Err(_) => return, + }; + let invalid = state + .lock() + .await + .validate_received_trusted_request(&request) + .is_err(); + if validated_tx.send(request).await.is_err() { + return; + } + if invalid { + return; + } + } + }); + ( + Box::pin(stream::unfold(validated_rx, |mut receiver| async move { + receiver.recv().await.map(|request| (request, receiver)) + })), + initial_requests_checked_rx, + ) +} + /// The only way to turn a wire-level freshness disposition into the internal /// [`InvocationFreshnessDisposition`]. Decoding and trust-sanitization are /// deliberately fused into a single function so that no gRPC entry point can @@ -53,8 +222,7 @@ fn sanitize_invocation_freshness_disposition( trusted_internal_caller: bool, ) -> InvocationFreshnessDisposition { let decoded = if wire_value - == golem_api_grpc::proto::golem::worker::v1::InvocationFreshnessDisposition::KnownFresh - as i32 + == golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh as i32 { InvocationFreshnessDisposition::KnownFresh } else { @@ -271,6 +439,123 @@ impl GrpcWorkerService for WorkerGrpcApi { })) } + type InvokeAgentSessionStream = InvocationResponseStream; + + async fn invoke_agent_session( + &self, + request: Request>, + ) -> Result, Status> { + let mut inbound = request.into_inner(); + let first = match inbound.message().await { + Ok(Some(request)) => request, + Ok(None) => { + return Ok(Response::new(service_failure_stream( + InvocationRejectionReason::Protocol, + "invocation request ended before start".to_string(), + None, + None, + ))); + } + Err(error) => { + return Err(error); + } + }; + let (idempotency_key, agent_id) = request_identity(&first); + let mut state = InvocationSessionState::default(); + if let Err(error) = state.validate_trusted_request(&first) { + return Ok(Response::new(service_failure_stream( + InvocationRejectionReason::Protocol, + error, + idempotency_key, + agent_id, + ))); + } + let mut start = match first + .request + .expect("validated invocation request has a payload") + { + invocation_request::Request::Start(start) => start, + invocation_request::Request::ResumeAttach(_) => { + return Ok(Response::new(service_failure_stream( + InvocationRejectionReason::ResumeUnsupported, + "resume-attach is not supported by live sessions".to_string(), + idempotency_key, + None, + ))); + } + _ => unreachable!("the session validator requires start or resume-attach first"), + }; + let auth = match start.auth_ctx.clone() { + Some(auth) => match auth.try_into() { + Ok(auth) => auth, + Err(error) => { + return Ok(Response::new(service_failure_stream( + InvocationRejectionReason::Validation, + format!("failed converting auth_ctx: {error}"), + idempotency_key, + agent_id, + ))); + } + }, + None => { + return Ok(Response::new(service_failure_stream( + InvocationRejectionReason::Validation, + "auth_ctx not found".to_string(), + idempotency_key, + agent_id, + ))); + } + }; + let trusted_internal_caller = matches!(&auth, AuthCtx::System | AuthCtx::Agent(_)); + start.freshness_disposition = match sanitize_invocation_freshness_disposition( + start.freshness_disposition, + trusted_internal_caller, + ) { + InvocationFreshnessDisposition::MayExist => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32 + } + InvocationFreshnessDisposition::KnownFresh => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh + as i32 + } + }; + let state = Arc::new(tokio::sync::Mutex::new(state)); + let (tail, initial_requests_checked) = validated_request_tail(inbound, state.clone()); + match self + .worker_service + .invoke_agent_session(start, tail, trusted_internal_caller, auth) + .await + { + Ok(response) => Ok(Response::new(validated_response_stream( + response, + state, + Some(initial_requests_checked), + ))), + Err(error) => { + let reason = match &error { + WorkerServiceError::AuthError(_) | WorkerServiceError::LimitError(_) => { + InvocationRejectionReason::Unauthorized + } + WorkerServiceError::ComponentNotFound(_) + | WorkerServiceError::AgentNotFound(_) + | WorkerServiceError::AccountIdNotFound(_) => { + InvocationRejectionReason::NotFound + } + WorkerServiceError::TypeChecker(_) => InvocationRejectionReason::Validation, + WorkerServiceError::InternalCallError(_) => InvocationRejectionReason::Internal, + _ => InvocationRejectionReason::Internal, + }; + Ok(Response::new(service_failure_stream( + reason, + error.to_string(), + idempotency_key, + agent_id, + ))) + } + } + } + async fn cancel_invocation( &self, request: Request, @@ -522,12 +807,6 @@ impl WorkerGrpcApi { .principal .unwrap_or_else(|| golem_common::model::agent::Principal::anonymous().into()); - let environment_id = request - .environment_id - .map(|id| id.try_into()) - .transpose() - .map_err(|e| bad_request_error(format!("invalid environment_id: {e}")))?; - let output = self .worker_service .invoke_agent( @@ -543,13 +822,16 @@ impl WorkerGrpcApi { config, auth, principal, - environment_id, ) .await?; let result_value = match &output.result { golem_common::model::AgentInvocationResult::AgentMethod { output } => { - Some(output.clone().into()) + Some(output.clone().try_into().map_err(|error| { + WorkerServiceError::Internal(format!( + "agent output cannot cross the gRPC boundary: {error}" + )) + })?) } _ => None, }; @@ -687,3 +969,206 @@ mod freshness_tests { ); } } + +#[cfg(test)] +mod protocol_tests { + use super::{validated_request_tail, validated_response_stream}; + use futures::{FutureExt, StreamExt, stream}; + use golem_api_grpc::invocation_session_protocol::InvocationSessionState; + use golem_api_grpc::proto::golem::common::Empty; + use golem_api_grpc::proto::golem::schema::{SchemaValue, schema_value}; + use golem_api_grpc::proto::golem::worker::{ + AgentId, IdempotencyKey, InvocationAccepted, InvocationRequest, InvocationResponse, + InvocationSessionCompletion, InvocationSessionResult, InvocationStart, invocation_request, + invocation_response, invocation_session_completion, invocation_session_result, + }; + use std::sync::Arc; + use test_r::test; + use tonic::Status; + + fn key() -> Option { + Some(IdempotencyKey { + value: "session-key".to_string(), + }) + } + + fn agent_id() -> Option { + Some(AgentId { + component_id: None, + name: "agent".to_string(), + }) + } + + fn start() -> InvocationRequest { + InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + input: Some(SchemaValue { + value: Some(schema_value::Value::U8Value(1)), + }), + idempotency_key: key(), + ..Default::default() + })), + } + } + + fn state_after_start() -> Arc> { + let mut state = InvocationSessionState::default(); + state.validate_trusted_request(&start()).unwrap(); + Arc::new(tokio::sync::Mutex::new(state)) + } + + fn response(response: invocation_response::Response) -> InvocationResponse { + InvocationResponse { + response: Some(response), + } + } + + fn accepted() -> InvocationResponse { + response(invocation_response::Response::Accepted( + InvocationAccepted { + agent_id: agent_id(), + idempotency_key: key(), + component_revision: Some(1), + }, + )) + } + + fn result() -> InvocationResponse { + response(invocation_response::Response::Result( + InvocationSessionResult { + result: Some(invocation_session_result::Result::NoResult(Empty {})), + component_revision: Some(1), + agent_id: agent_id(), + idempotency_key: key(), + ..Default::default() + }, + )) + } + + fn successful_completion() -> InvocationResponse { + response(invocation_response::Response::Finished( + InvocationSessionCompletion { + outcome: Some(invocation_session_completion::Outcome::Success(Empty {})), + }, + )) + } + + #[test] + async fn request_transport_error_closes_the_internal_request_stream() { + let inbound = stream::iter([Err(Status::unavailable("request transport failed"))]); + let (mut tail, initial_requests_checked) = + validated_request_tail(inbound, state_after_start()); + + initial_requests_checked.await.unwrap(); + assert!(tail.next().await.is_none()); + } + + #[test] + async fn malformed_request_tail_is_forwarded_for_semantic_terminalization() { + let malformed = start(); + let inbound = stream::iter([Ok::<_, Status>(malformed.clone())]); + let (mut tail, initial_requests_checked) = + validated_request_tail(inbound, state_after_start()); + + initial_requests_checked.await.unwrap(); + assert_eq!(tail.next().await, Some(malformed)); + assert!(tail.next().await.is_none()); + } + + #[test] + async fn response_transport_error_is_preserved() { + let inbound = stream::iter([Err(Status::unavailable("response transport failed"))]); + let mut responses = validated_response_stream(inbound, state_after_start(), None); + + let error = responses.next().await.unwrap().unwrap_err(); + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(error.message().contains("response transport failed")); + assert!(responses.next().await.is_none()); + } + + #[test] + async fn valid_response_preserves_stream_free_lifecycle() { + let inbound = stream::iter([ + Ok::<_, Status>(accepted()), + Ok(result()), + Ok(successful_completion()), + ]); + let mut responses = validated_response_stream(inbound, state_after_start(), None); + + assert!(matches!( + responses.next().await.unwrap().unwrap().response, + Some(invocation_response::Response::Accepted(_)) + )); + assert!(matches!( + responses.next().await.unwrap().unwrap().response, + Some(invocation_response::Response::Result(_)) + )); + assert!(matches!( + responses.next().await.unwrap().unwrap().response, + Some(invocation_response::Response::Finished(_)) + )); + assert!(responses.next().await.is_none()); + } + + #[test] + async fn terminal_response_is_withheld_until_the_upstream_closes_cleanly() { + let (sender, receiver) = tokio::sync::mpsc::channel(4); + sender.send(Ok::<_, Status>(accepted())).await.unwrap(); + sender.send(Ok(result())).await.unwrap(); + sender.send(Ok(successful_completion())).await.unwrap(); + let mut responses = validated_response_stream( + tokio_stream::wrappers::ReceiverStream::new(receiver), + state_after_start(), + None, + ); + + assert!(responses.next().await.unwrap().is_ok()); + assert!(responses.next().await.unwrap().is_ok()); + assert!(responses.next().now_or_never().is_none()); + + drop(sender); + assert!(matches!( + responses.next().await.unwrap().unwrap().response, + Some(invocation_response::Response::Finished(_)) + )); + assert!(responses.next().await.is_none()); + } + + #[test] + async fn response_error_after_terminal_suppresses_the_terminal() { + let inbound = stream::iter([ + Ok::<_, Status>(accepted()), + Ok(result()), + Ok(successful_completion()), + Err(Status::unavailable( + "response transport failed after terminal", + )), + ]); + let mut responses = validated_response_stream(inbound, state_after_start(), None); + + assert!(responses.next().await.unwrap().is_ok()); + assert!(responses.next().await.unwrap().is_ok()); + let error = responses.next().await.unwrap().unwrap_err(); + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(error.message().contains("after terminal")); + assert!(responses.next().await.is_none()); + } + + #[test] + async fn repeated_response_terminal_is_a_protocol_status() { + let inbound = stream::iter([ + Ok::<_, Status>(accepted()), + Ok(result()), + Ok(successful_completion()), + Ok(successful_completion()), + ]); + let mut responses = validated_response_stream(inbound, state_after_start(), None); + + assert!(responses.next().await.unwrap().is_ok()); + assert!(responses.next().await.unwrap().is_ok()); + let error = responses.next().await.unwrap().unwrap_err(); + assert_eq!(error.code(), tonic::Code::Internal); + assert!(error.message().contains("after completion")); + assert!(responses.next().await.is_none()); + } +} diff --git a/golem-worker-service/src/mcp/invoke/resource.rs b/golem-worker-service/src/mcp/invoke/resource.rs index e7bb29913d..033afa73e3 100644 --- a/golem-worker-service/src/mcp/invoke/resource.rs +++ b/golem-worker-service/src/mcp/invoke/resource.rs @@ -82,7 +82,9 @@ pub async fn invoke_resource( let method_parameters = SchemaValue::Record { fields: vec![] }; let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - method_parameters.into(); + method_parameters.try_into().map_err(|error| { + ErrorData::internal_error(format!("Failed to encode method parameters: {error}"), None) + })?; let principal = Principal::anonymous(); let proto_principal: golem_api_grpc::proto::golem::component::Principal = principal.into(); @@ -108,7 +110,6 @@ pub async fn invoke_resource( Vec::new(), auth_ctx, proto_principal, - None, ) .await .map_err(|e| { @@ -571,6 +572,21 @@ mod tests { #[test] async fn invoke_resource_auto_generates_phantom_for_ephemeral_agents() { + let constructor = AgentConstructorSchema { + name: None, + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![]), + }; + let method = AgentMethodSchema { + name: "read".to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![]), + output_schema: OutputSchema::Unit, + http_endpoint: vec![], + read_only: None, + }; let harness = InvocationHarness::new_with_agent_mode( AgentInvocationOutput { result: golem_common::model::AgentInvocationResult::AgentInitialization, @@ -583,6 +599,8 @@ mod tests { agent_fingerprint: None, }, AgentMode::Ephemeral, + constructor.clone(), + vec![method.clone()], ); let resource = AgentMcpResource { kind: AgentMcpResourceKind::Static(Annotated::new( @@ -602,21 +620,8 @@ mod tests { account_id: harness.account_id, schema_graph: Arc::new(SchemaGraph::empty()), account_email: golem_common::model::account::AccountEmail::new("mcp@golem"), - constructor: AgentConstructorSchema { - name: None, - description: String::new(), - prompt_hint: None, - input_schema: InputSchema::Parameters(vec![]), - }, - method: AgentMethodSchema { - name: "read".to_string(), - description: String::new(), - prompt_hint: None, - input_schema: InputSchema::Parameters(vec![]), - output_schema: OutputSchema::Unit, - http_endpoint: vec![], - read_only: None, - }, + constructor, + method, component_id: harness.component_id, agent_type_name: AgentTypeName("mcp-agent".to_string()), agent_mode: AgentMode::Ephemeral, diff --git a/golem-worker-service/src/mcp/invoke/test_support.rs b/golem-worker-service/src/mcp/invoke/test_support.rs index 0809bb3f58..c3e6f4759d 100644 --- a/golem-worker-service/src/mcp/invoke/test_support.rs +++ b/golem-worker-service/src/mcp/invoke/test_support.rs @@ -16,7 +16,9 @@ use crate::service::agent_resolution_cache::AgentResolutionCache; use crate::service::auth::{AuthService, AuthServiceError}; use crate::service::component::{ComponentService, ComponentServiceError}; use crate::service::limit::{LimitService, LimitServiceError}; -use crate::service::worker::{WorkerClient, WorkerResult, WorkerService, WorkerStream}; +use crate::service::worker::{ + WorkerClient, WorkerResult, WorkerService, WorkerServiceError, WorkerStream, +}; use async_trait::async_trait; use bytes::Bytes; use chrono::Utc; @@ -37,7 +39,7 @@ use golem_common::model::environment::{EnvironmentId, EnvironmentName}; use golem_common::model::oplog::{OplogCursor, OplogIndex}; use golem_common::model::worker::{AgentConfigEntryDto, AgentMetadataDto, RevertWorkerTarget}; use golem_common::model::{AgentFilter, AgentFingerprint, AgentId, IdempotencyKey, ScanCursor}; -use golem_common::schema::{AgentConstructorSchema, AgentTypeSchema, InputSchema, SchemaGraph}; +use golem_common::schema::{AgentConstructorSchema, AgentTypeSchema, SchemaGraph}; use golem_service_base::clients::registry::{RegistryService, RegistryServiceError}; use golem_service_base::model::auth::AuthCtx; use golem_service_base::model::component::Component; @@ -351,11 +353,11 @@ impl WorkerClient for RecordingWorkerClient { async fn get_metadata( &self, - _: &AgentId, + agent_id: &AgentId, _: EnvironmentId, _: AuthCtx, ) -> WorkerResult { - unimplemented!() + Err(WorkerServiceError::AgentNotFound(agent_id.clone())) } async fn find_metadata( @@ -545,13 +547,19 @@ pub(crate) struct InvocationHarness { } impl InvocationHarness { - pub(crate) fn new(invocation_output: AgentInvocationOutput) -> Self { - Self::new_with_agent_mode(invocation_output, AgentMode::Durable) + pub(crate) fn new( + invocation_output: AgentInvocationOutput, + constructor: AgentConstructorSchema, + methods: Vec, + ) -> Self { + Self::new_with_agent_mode(invocation_output, AgentMode::Durable, constructor, methods) } pub(crate) fn new_with_agent_mode( invocation_output: AgentInvocationOutput, agent_mode: AgentMode, + constructor: AgentConstructorSchema, + methods: Vec, ) -> Self { let component_id = ComponentId(Uuid::new_v4()); let environment_id = EnvironmentId(Uuid::new_v4()); @@ -579,13 +587,8 @@ impl InvocationHarness { 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![]), - }, - methods: vec![], + constructor, + methods, dependencies: vec![], mode: agent_mode, http_mount: None, diff --git a/golem-worker-service/src/mcp/invoke/tool.rs b/golem-worker-service/src/mcp/invoke/tool.rs index 40f414eef0..d5844e982e 100644 --- a/golem-worker-service/src/mcp/invoke/tool.rs +++ b/golem-worker-service/src/mcp/invoke/tool.rs @@ -91,7 +91,9 @@ pub async fn invoke_tool( })?; let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - method_parameters.into(); + method_parameters.try_into().map_err(|error| { + ErrorData::internal_error(format!("Failed to encode method parameters: {error}"), None) + })?; let principal = Principal::anonymous(); let proto_principal: golem_api_grpc::proto::golem::component::Principal = principal.into(); @@ -117,7 +119,6 @@ pub async fn invoke_tool( Vec::new(), auth_ctx, proto_principal, - None, ) .await .map_err(|e| { @@ -696,6 +697,21 @@ mod tests { #[test] async fn invoke_tool_auto_generates_phantom_for_ephemeral_agents() { + let constructor = AgentConstructorSchema { + name: None, + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![]), + }; + let method = AgentMethodSchema { + name: "run".to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![]), + output_schema: OutputSchema::Unit, + http_endpoint: vec![], + read_only: None, + }; let harness = InvocationHarness::new_with_agent_mode( AgentInvocationOutput { result: golem_common::model::AgentInvocationResult::AgentInitialization, @@ -708,6 +724,8 @@ mod tests { agent_fingerprint: None, }, AgentMode::Ephemeral, + constructor.clone(), + vec![method.clone()], ); let tool = AgentMcpTool { tool: Tool { @@ -725,21 +743,8 @@ mod tests { account_id: harness.account_id, schema_graph: Arc::new(SchemaGraph::empty()), account_email: golem_common::model::account::AccountEmail::new("mcp@golem"), - constructor: AgentConstructorSchema { - name: None, - description: String::new(), - prompt_hint: None, - input_schema: InputSchema::Parameters(vec![]), - }, - method: AgentMethodSchema { - name: "run".to_string(), - description: String::new(), - prompt_hint: None, - input_schema: InputSchema::Parameters(vec![]), - output_schema: OutputSchema::Unit, - http_endpoint: vec![], - read_only: None, - }, + constructor, + method, component_id: harness.component_id, agent_type_name: AgentTypeName("mcp-agent".to_string()), agent_mode: AgentMode::Ephemeral, @@ -760,18 +765,43 @@ mod tests { // `method_id`. The invoke path must translate those advertised names // back and route each value to the correct side (different types make // a swap observable: constructor = string, method = u32). - let harness = InvocationHarness::new(AgentInvocationOutput { - result: golem_common::model::AgentInvocationResult::AgentMethod { - output: SchemaValue::Tuple { elements: vec![] }, + let constructor = AgentConstructorSchema { + name: None, + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![NamedField::user_supplied( + "id", + SchemaType::string(), + )]), + }; + let method = AgentMethodSchema { + name: "run".to_string(), + description: String::new(), + prompt_hint: None, + input_schema: InputSchema::Parameters(vec![NamedField::user_supplied( + "id", + SchemaType::u32(), + )]), + output_schema: OutputSchema::Unit, + http_endpoint: vec![], + read_only: None, + }; + let harness = InvocationHarness::new( + AgentInvocationOutput { + result: golem_common::model::AgentInvocationResult::AgentMethod { + output: SchemaValue::Tuple { elements: vec![] }, + }, + consumed_fuel: None, + invocation_status: None, + component_revision: None, + agent_id: None, + idempotency_key: None, + oplog_index: None, + agent_fingerprint: None, }, - consumed_fuel: None, - invocation_status: None, - component_revision: None, - agent_id: None, - idempotency_key: None, - oplog_index: None, - agent_fingerprint: None, - }); + constructor.clone(), + vec![method.clone()], + ); let tool = AgentMcpTool { tool: Tool { name: Cow::Borrowed("mcp-agent-run"), @@ -788,27 +818,8 @@ mod tests { account_id: harness.account_id, account_email: harness.account_email.clone(), schema_graph: Arc::new(SchemaGraph::empty()), - constructor: AgentConstructorSchema { - name: None, - description: String::new(), - prompt_hint: None, - input_schema: InputSchema::Parameters(vec![NamedField::user_supplied( - "id", - SchemaType::string(), - )]), - }, - method: AgentMethodSchema { - name: "run".to_string(), - description: String::new(), - prompt_hint: None, - input_schema: InputSchema::Parameters(vec![NamedField::user_supplied( - "id", - SchemaType::u32(), - )]), - output_schema: OutputSchema::Unit, - http_endpoint: vec![], - read_only: None, - }, + constructor, + method, component_id: harness.component_id, agent_type_name: AgentTypeName("mcp-agent".to_string()), agent_mode: AgentMode::Ephemeral, diff --git a/golem-worker-service/src/service/worker/client.rs b/golem-worker-service/src/service/worker/client.rs index 697d26b05f..78e76d2ea3 100644 --- a/golem-worker-service/src/service/worker/client.rs +++ b/golem-worker-service/src/service/worker/client.rs @@ -17,11 +17,21 @@ use super::{ AllExecutors, CallWorkerExecutorError, HasWorkerExecutorClients, RandomExecutor, ResponseMapResult, RoutingLogic, WorkerServiceError, WorkerStream, }; +use crate::service::auth::AuthServiceError; use async_trait::async_trait; use bytes::Bytes; use futures::stream::TryStreamExt; use futures::{Stream, StreamExt}; -use golem_api_grpc::proto::golem::worker::{InvocationContext, LogEvent}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::worker::invocation_request; +use golem_api_grpc::proto::golem::worker::invocation_response; +use golem_api_grpc::proto::golem::worker::invocation_session_completion; +use golem_api_grpc::proto::golem::worker::invocation_session_result; +use golem_api_grpc::proto::golem::worker::{ + InvocationContext, InvocationFailure, InvocationFailureKind, InvocationRejected, + InvocationRejectionReason, InvocationRequest, InvocationResponse, InvocationSessionResult, + InvocationStart, LogEvent, +}; use golem_api_grpc::proto::golem::workerexecutor; use golem_api_grpc::proto::golem::workerexecutor::v1::worker_executor_client::WorkerExecutorClient; use golem_api_grpc::proto::golem::workerexecutor::v1::{ @@ -53,11 +63,14 @@ use golem_service_base::grpc::client::MultiTargetGrpcClient; use golem_service_base::model::auth::AuthCtx; use golem_service_base::model::{ComponentFileSystemNode, GetOplogResponse}; use golem_service_base::service::routing_table::{HasRoutingTableService, RoutingTableService}; +use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::{collections::HashMap, sync::Arc}; -use tonic::Code; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; use tonic::transport::Channel; +use tonic::{Code, Status}; use tonic_tracing_opentelemetry::middleware::client::OtelGrpcService; fn freshness_disposition_for_dispatch( @@ -73,6 +86,203 @@ fn freshness_disposition_for_dispatch( } } +pub type InvocationRequestStream = Pin + Send + 'static>>; +pub type InvocationResponseStream = + Pin> + Send + 'static>>; +type InvocationSessionCall<'a> = Pin< + Box< + dyn Future>, Status>> + + Send + + 'a, + >, +>; + +fn invoke_agent_session_once<'a>( + client: &'a mut WorkerExecutorClient>, + request: Option, +) -> InvocationSessionCall<'a> { + match request { + Some(request) => Box::pin(client.invoke_agent_session(request)), + None => Box::pin(std::future::ready(Err(Status::aborted( + "invocation session request was already consumed", + )))), + } +} + +#[derive(Debug)] +enum OneShotInvocationSessionResult { + Success(AgentInvocationOutput), + Rejected(InvocationRejected), + Failure(InvocationFailure), + ProtocolFailure(String), +} + +fn protocol_failure(details: impl Into) -> OneShotInvocationSessionResult { + OneShotInvocationSessionResult::ProtocolFailure(details.into()) +} + +fn decode_invocation_rejection(rejected: InvocationRejected) -> WorkerServiceError { + match InvocationRejectionReason::try_from(rejected.reason) { + Ok(InvocationRejectionReason::NotFound) => rejected + .agent_id + .map(TryInto::try_into) + .transpose() + .map_err(WorkerServiceError::TypeChecker) + .and_then(|agent_id| { + agent_id + .map(WorkerServiceError::AgentNotFound) + .ok_or_else(|| WorkerServiceError::Internal(rejected.error.clone())) + }) + .unwrap_or_else(|error| error), + Ok(InvocationRejectionReason::Unauthorized) => { + WorkerServiceError::AuthError(AuthServiceError::CouldNotAuthenticate) + } + Ok(InvocationRejectionReason::Internal) => WorkerServiceError::Internal(rejected.error), + _ => WorkerServiceError::TypeChecker(rejected.error), + } +} + +fn decode_invocation_failure(failure: InvocationFailure) -> WorkerExecutorError { + if failure.kind == InvocationFailureKind::Protocol as i32 { + WorkerExecutorError::invalid_request(failure.message) + } else if let Some(worker_error) = failure.worker_error { + worker_error + .try_into() + .unwrap_or_else(|error| WorkerExecutorError::Unknown { + details: format!("failed to decode worker execution error: {error}"), + }) + } else { + WorkerExecutorError::Unknown { + details: failure.message, + } + } +} + +fn decode_invocation_result( + wire: InvocationSessionResult, +) -> Result { + let result = match wire.result.ok_or("invocation result has no payload")? { + invocation_session_result::Result::MethodResult(value) => { + AgentInvocationResult::AgentMethod { + output: value.try_into()?, + } + } + invocation_session_result::Result::NoResult(_) => { + AgentInvocationResult::AgentInitialization + } + }; + let invocation_status = wire.status.and_then(|status| { + golem_api_grpc::proto::golem::worker::InvocationStatus::try_from(status) + .ok() + .map(InvocationStatus::from) + }); + Ok(AgentInvocationOutput { + result, + consumed_fuel: wire.fuel_consumed, + invocation_status, + component_revision: wire + .component_revision + .map(ComponentRevision::new) + .transpose() + .map_err(|error| error.to_string())?, + agent_id: wire.agent_id.map(TryInto::try_into).transpose()?, + idempotency_key: wire.idempotency_key.map(Into::into), + oplog_index: wire.oplog_index.map(OplogIndex::from_u64), + agent_fingerprint: wire + .agent_fingerprint + .map(|uuid| AgentFingerprint(uuid.into())), + }) +} + +async fn run_one_shot_invocation_session( + client: &mut WorkerExecutorClient>, + start: InvocationStart, +) -> Result { + let (requests, receiver) = mpsc::channel(1); + let request = InvocationRequest { + request: Some(invocation_request::Request::Start(start)), + }; + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&request) + .map_err(Status::invalid_argument)?; + requests + .send(request) + .await + .map_err(|_| Status::unavailable("invocation session request ended before start"))?; + let responses = client + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + collect_one_shot_invocation_session(responses, state).await +} + +async fn collect_one_shot_invocation_session( + mut responses: S, + mut state: InvocationSessionState, +) -> Result +where + S: Stream> + Unpin, +{ + let mut result = None; + let mut terminal_outcome = None; + + while let Some(response) = responses.next().await.transpose()? { + if let Err(details) = state.validate_response(&response) { + return Ok(protocol_failure(details)); + } + match response.response { + Some(invocation_response::Response::Accepted(_)) => {} + Some(invocation_response::Response::Rejected(rejected)) => { + terminal_outcome = Some(Ok(OneShotInvocationSessionResult::Rejected(rejected))); + } + Some(invocation_response::Response::Result(invocation_result)) => { + result = match decode_invocation_result(invocation_result) { + Ok(result) => Some(result), + Err(details) => return Ok(protocol_failure(details)), + }; + } + Some(invocation_response::Response::Finished(finished)) => { + terminal_outcome = Some(match finished.outcome { + Some(invocation_session_completion::Outcome::Success(_)) => result + .take() + .map(OneShotInvocationSessionResult::Success) + .ok_or_else(|| Status::internal("invocation completed without a result")), + Some(invocation_session_completion::Outcome::Failure(failure)) => { + if failure.kind == InvocationFailureKind::Transport as i32 { + Err(Status::unavailable(failure.message)) + } else { + Ok(OneShotInvocationSessionResult::Failure(failure)) + } + } + None => Ok(protocol_failure("invocation completion has no outcome")), + }); + } + Some( + invocation_response::Response::OutputItem(_) + | invocation_response::Response::OutputEnd(_) + | invocation_response::Response::OutputError(_) + | invocation_response::Response::InputAck(_) + | invocation_response::Response::StreamCancel(_), + ) => { + return Ok(protocol_failure( + "a non-streaming invocation received a stream frame", + )); + } + Some(invocation_response::Response::AttachmentRevoked(_)) => { + unreachable!("response validation rejects attachment revocation") + } + None => unreachable!("response validation rejects empty frames"), + } + } + + terminal_outcome.unwrap_or_else(|| { + Err(Status::unavailable( + "invocation session response ended before completion", + )) + }) +} + #[async_trait] pub trait WorkerClient: Send + Sync { async fn create( @@ -262,6 +472,16 @@ pub trait WorkerClient: Send + Sync { principal: golem_api_grpc::proto::golem::component::Principal, ) -> WorkerResult; + async fn invoke_agent_session( + &self, + _agent_id: &AgentId, + _request: InvocationRequestStream, + ) -> WorkerResult { + Err(WorkerServiceError::Internal( + "invocation sessions are not supported by this worker client".to_string(), + )) + } + async fn process_oplog_entries( &self, target_agent_id: &AgentId, @@ -1391,89 +1611,50 @@ impl WorkerClient for WorkerExecutorWorkerClient { let result = self .call_worker_executor( agent_id.clone(), - "invoke_agent", + "invoke_agent_session", move |worker_executor_client| { let freshness_disposition = freshness_disposition_for_dispatch(freshness_disposition, &first_dispatch); - Box::pin(worker_executor_client.invoke_agent( - workerexecutor::v1::InvokeAgentRequest { - agent_id: Some(agent_id_clone.clone().into()), - method_name: method_name.clone(), - method_parameters: method_parameters.clone(), - mode, - schedule_at, - idempotency_key: idempotency_key.clone().map(|k| k.into()), - component_owner_account_id: Some(account_id.into()), - environment_id: Some(environment_id.into()), - auth_ctx: Some(auth_ctx.clone().into()), - context: invocation_context.clone(), - principal: Some(principal.clone()), - freshness_disposition: match freshness_disposition { - InvocationFreshnessDisposition::MayExist => { - workerexecutor::v1::InvocationFreshnessDisposition::MayExist - as i32 - } - InvocationFreshnessDisposition::KnownFresh => { - workerexecutor::v1::InvocationFreshnessDisposition::KnownFresh - as i32 - } - }, - config: config.clone().into_iter().map(Into::into).collect(), + let start = InvocationStart { + agent_id: Some(agent_id_clone.clone().into()), + method_name: method_name.clone(), + input: method_parameters.clone(), + idempotency_key: idempotency_key.clone().map(Into::into), + context: invocation_context.clone(), + auth_ctx: Some(auth_ctx.clone().into()), + principal: Some(principal.clone()), + environment_id: Some(environment_id.into()), + config: config.clone().into_iter().map(Into::into).collect(), + component_owner_account_id: Some(account_id.into()), + mode, + schedule_at, + freshness_disposition: match freshness_disposition { + InvocationFreshnessDisposition::MayExist => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32 + } + InvocationFreshnessDisposition::KnownFresh => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh + as i32 + } }, + }; + Box::pin(run_one_shot_invocation_session( + worker_executor_client, + start, )) }, - |response| match response.into_inner() { - workerexecutor::v1::InvokeAgentResponse { - result: - Some(workerexecutor::v1::invoke_agent_response::Result::Success( - workerexecutor::v1::InvokeAgentSuccess { - result, - fuel_consumed, - component_revision, - status, - oplog_index, - agent_fingerprint, - agent_id, - idempotency_key, - }, - )), - } => { - let invocation_result = match result { - Some(proto_val) => { - let output = golem_common::schema::SchemaValue::try_from(proto_val) - .map_err(WorkerExecutorError::unknown)?; - AgentInvocationResult::AgentMethod { output } - } - None => AgentInvocationResult::AgentInitialization, - }; - let invocation_status = status.and_then(|s| { - golem_api_grpc::proto::golem::worker::InvocationStatus::try_from(s) - .ok() - .map(InvocationStatus::from) - }); - Ok(AgentInvocationOutput { - result: invocation_result, - consumed_fuel: fuel_consumed, - invocation_status, - component_revision: component_revision - .map(ComponentRevision::new) - .transpose() - .map_err(|err| WorkerExecutorError::unknown(err.to_string()))?, - agent_id: agent_id - .map(TryInto::try_into) - .transpose() - .map_err(|err: String| WorkerExecutorError::unknown(err))?, - idempotency_key: idempotency_key.map(Into::into), - oplog_index: oplog_index.map(OplogIndex::from_u64), - agent_fingerprint: agent_fingerprint - .map(|uuid| AgentFingerprint(uuid.into())), - }) + |outcome| match outcome { + OneShotInvocationSessionResult::Success(output) => Ok(output), + OneShotInvocationSessionResult::Rejected(rejected) => { + Err(decode_invocation_rejection(rejected).into()) + } + OneShotInvocationSessionResult::Failure(failure) => { + Err(decode_invocation_failure(failure).into()) + } + OneShotInvocationSessionResult::ProtocolFailure(details) => { + Err(WorkerExecutorError::invalid_request(details).into()) } - workerexecutor::v1::InvokeAgentResponse { - result: - Some(workerexecutor::v1::invoke_agent_response::Result::Failure(err)), - } => Err(err.into()), - workerexecutor::v1::InvokeAgentResponse { .. } => Err("Empty response".into()), }, WorkerServiceError::InternalCallError, ) @@ -1482,6 +1663,48 @@ impl WorkerClient for WorkerExecutorWorkerClient { Ok(result) } + async fn invoke_agent_session( + &self, + agent_id: &AgentId, + request: InvocationRequestStream, + ) -> WorkerResult { + let routing_table = self + .routing_table_service + .get_routing_table() + .await + .map_err(|error| { + WorkerServiceError::InternalCallError( + CallWorkerExecutorError::FailedToGetRoutingTable(error), + ) + })?; + let pod = routing_table.lookup(agent_id).ok_or_else(|| { + WorkerServiceError::InternalCallError(CallWorkerExecutorError::FailedToConnectToPod( + Status::unavailable(format!("no active shard for agent {agent_id}")), + )) + })?; + let request = Arc::new(std::sync::Mutex::new(Some(request))); + let response = self + .worker_executor_clients + .call_without_retry( + "invoke_agent_session", + pod.uri(self.worker_executor_clients.uses_tls()), + move |worker_executor_client| { + let request = request + .lock() + .unwrap_or_else(|poison| poison.into_inner()) + .take(); + invoke_agent_session_once(worker_executor_client, request) + }, + ) + .await + .map_err(|status| { + WorkerServiceError::InternalCallError( + CallWorkerExecutorError::FailedToConnectToPod(status), + ) + })?; + Ok(Box::pin(response.into_inner())) + } + async fn process_oplog_entries( &self, target_agent_id: &AgentId, @@ -1600,3 +1823,600 @@ mod freshness_tests { ); } } + +#[cfg(test)] +mod one_shot_session_tests { + use super::{ + OneShotInvocationSessionResult, collect_one_shot_invocation_session, + decode_invocation_failure, + }; + use futures::stream; + use golem_api_grpc::invocation_session_protocol::InvocationSessionState; + use golem_api_grpc::proto::golem::common::Empty; + use golem_api_grpc::proto::golem::schema::{SchemaValue, schema_value}; + use golem_api_grpc::proto::golem::worker::{ + AgentId, IdempotencyKey, InvocationAccepted, InvocationFailure, InvocationFailureKind, + InvocationRequest, InvocationResponse, InvocationSessionCompletion, + InvocationSessionResult, InvocationStart, invocation_request, invocation_response, + invocation_session_completion, invocation_session_result, + }; + use golem_common::model::AgentFingerprint; + use golem_common::model::oplog::{AgentError, OplogIndex}; + use golem_service_base::error::worker_executor::WorkerExecutorError; + use test_r::test; + use tonic::Status; + + fn key() -> Option { + Some(IdempotencyKey { + value: "session-key".to_string(), + }) + } + + fn agent_id() -> Option { + Some( + golem_common::model::AgentId { + component_id: golem_common::model::component::ComponentId(uuid::Uuid::nil()), + agent_id: "agent".to_string(), + } + .into(), + ) + } + + fn state_after_start() -> InvocationSessionState { + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + input: Some(SchemaValue { + value: Some(schema_value::Value::U8Value(1)), + }), + idempotency_key: key(), + ..Default::default() + })), + }) + .unwrap(); + state + } + + fn frame(response: invocation_response::Response) -> Result { + Ok(InvocationResponse { + response: Some(response), + }) + } + + fn accepted() -> invocation_response::Response { + invocation_response::Response::Accepted(InvocationAccepted { + agent_id: agent_id(), + idempotency_key: key(), + component_revision: Some(3), + }) + } + + fn no_result() -> invocation_response::Response { + invocation_response::Response::Result(InvocationSessionResult { + result: Some(invocation_session_result::Result::NoResult(Empty {})), + component_revision: Some(3), + agent_id: agent_id(), + idempotency_key: key(), + ..Default::default() + }) + } + + fn finished(outcome: invocation_session_completion::Outcome) -> invocation_response::Response { + invocation_response::Response::Finished(InvocationSessionCompletion { + outcome: Some(outcome), + }) + } + + #[test] + async fn successful_session_preserves_unary_result_metadata() { + let responses = stream::iter([ + frame(accepted()), + frame(invocation_response::Response::Result( + InvocationSessionResult { + result: Some(invocation_session_result::Result::NoResult(Empty {})), + component_revision: Some(3), + agent_id: agent_id(), + idempotency_key: key(), + fuel_consumed: Some(17), + status: Some( + golem_api_grpc::proto::golem::worker::InvocationStatus::Complete as i32, + ), + oplog_index: Some(29), + agent_fingerprint: Some(uuid::Uuid::nil().into()), + }, + )), + frame(finished(invocation_session_completion::Outcome::Success( + Empty {}, + ))), + ]); + + let result = collect_one_shot_invocation_session(responses, state_after_start()) + .await + .unwrap(); + let OneShotInvocationSessionResult::Success(output) = result else { + panic!("expected a successful invocation session"); + }; + assert_eq!(output.agent_id.unwrap().agent_id, "agent"); + assert_eq!(output.idempotency_key.unwrap().value, "session-key"); + assert_eq!(output.component_revision.unwrap().get(), 3); + assert_eq!(output.consumed_fuel, Some(17)); + assert_eq!( + output.invocation_status, + Some(golem_common::model::InvocationStatus::Complete) + ); + assert_eq!(output.oplog_index, Some(OplogIndex::from_u64(29))); + assert_eq!( + output.agent_fingerprint, + Some(AgentFingerprint(uuid::Uuid::nil())) + ); + } + + #[test] + async fn typed_executor_failure_survives_the_session_adapter() { + let expected = WorkerExecutorError::InvocationFailed { + error: AgentError::Unknown("failed".to_string()), + stderr: "guest stderr".to_string(), + }; + let responses = stream::iter([ + frame(accepted()), + frame(finished(invocation_session_completion::Outcome::Failure( + InvocationFailure { + kind: InvocationFailureKind::Execution as i32, + code: "execution".to_string(), + message: "failed".to_string(), + worker_error: Some(expected.clone().into()), + }, + ))), + ]); + + let result = collect_one_shot_invocation_session(responses, state_after_start()) + .await + .unwrap(); + let OneShotInvocationSessionResult::Failure(failure) = result else { + panic!("expected an invocation failure"); + }; + assert_eq!(decode_invocation_failure(failure), expected); + } + + #[test] + async fn successful_completion_before_result_is_a_protocol_failure() { + let responses = stream::iter([ + frame(accepted()), + frame(finished(invocation_session_completion::Outcome::Success( + Empty {}, + ))), + ]); + + let result = collect_one_shot_invocation_session(responses, state_after_start()) + .await + .unwrap(); + assert!(matches!( + result, + OneShotInvocationSessionResult::ProtocolFailure(details) + if details.contains("before publishing a result") + )); + } + + #[test] + async fn session_is_drained_and_rejects_frames_after_completion() { + let responses = stream::iter([ + frame(accepted()), + frame(no_result()), + frame(finished(invocation_session_completion::Outcome::Success( + Empty {}, + ))), + frame(no_result()), + ]); + + let result = collect_one_shot_invocation_session(responses, state_after_start()) + .await + .unwrap(); + assert!(matches!( + result, + OneShotInvocationSessionResult::ProtocolFailure(details) + if details.contains("after completion") + )); + } + + #[test] + async fn response_transport_error_after_completion_is_retriable() { + let responses = stream::iter([ + frame(accepted()), + frame(no_result()), + frame(finished(invocation_session_completion::Outcome::Success( + Empty {}, + ))), + Err(Status::unavailable("response did not close cleanly")), + ]); + + let error = collect_one_shot_invocation_session(responses, state_after_start()) + .await + .unwrap_err(); + assert_eq!(error.code(), tonic::Code::Unavailable); + } + + #[test] + async fn typed_transport_failure_is_retriable() { + let responses = stream::iter([ + frame(accepted()), + frame(finished(invocation_session_completion::Outcome::Failure( + InvocationFailure { + kind: InvocationFailureKind::Transport as i32, + code: "transport".to_string(), + message: "request transport failed".to_string(), + worker_error: None, + }, + ))), + ]); + + let error = collect_one_shot_invocation_session(responses, state_after_start()) + .await + .unwrap_err(); + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(error.message().contains("request transport failed")); + } +} + +#[cfg(test)] +mod rejection_mapping_tests { + use super::{WorkerClient, WorkerExecutorWorkerClient, decode_invocation_rejection}; + use futures::{Stream, stream}; + use golem_api_grpc::proto::golem::schema::{SchemaValue, schema_value}; + use golem_api_grpc::proto::golem::shardmanager::{ + IpAddress, Pod as GrpcPod, RoutingTable as GrpcRoutingTable, RoutingTableEntry, ShardId, + ip_address, + }; + use golem_api_grpc::proto::golem::worker::v1::{AgentError, agent_error}; + use golem_api_grpc::proto::golem::worker::{ + InvocationRejected, InvocationRejectionReason, InvocationRequest, InvocationResponse, + invocation_response, + }; + use golem_api_grpc::proto::golem::workerexecutor::v1::worker_executor_server::{ + WorkerExecutor, WorkerExecutorServer, + }; + use golem_api_grpc::proto::golem::workerexecutor::v1::*; + use golem_common::model::account::AccountId; + use golem_common::model::agent::{InvocationFreshnessDisposition, Principal}; + use golem_common::model::component::ComponentId; + use golem_common::model::environment::EnvironmentId; + use golem_common::model::quota::{ResourceDefinitionId, ResourceName}; + use golem_common::model::{AgentId, RetryConfig, RoutingTable}; + use golem_service_base::clients::shard_manager::{ + BatchRenewalEntry, QuotaError, ShardManager, ShardManagerError, + }; + use golem_service_base::grpc::client::{GrpcClientConfig, MultiTargetGrpcClient}; + use golem_service_base::model::auth::AuthCtx; + use golem_service_base::model::quota_lease::{PendingReservation, QuotaLease}; + use golem_service_base::service::routing_table::{RoutingTableConfig, RoutingTableService}; + use std::net::Ipv4Addr; + use std::pin::Pin; + use std::sync::Arc; + use test_r::test; + use tokio::net::TcpListener; + use tokio_stream::wrappers::TcpListenerStream; + use tonic::codec::CompressionEncoding; + use tonic::{Request, Response, Status}; + use tonic_tracing_opentelemetry::middleware::client::OtelGrpcService; + + fn public_error_for_rejection(reason: InvocationRejectionReason) -> agent_error::Error { + let rejection = InvocationRejected { + reason: reason as i32, + error: "rejected".to_string(), + ..Default::default() + }; + let error: AgentError = decode_invocation_rejection(rejection).into(); + error.error.expect("missing public error") + } + + #[test] + fn validation_rejection_remains_a_bad_request() { + assert!(matches!( + public_error_for_rejection(InvocationRejectionReason::Validation), + agent_error::Error::BadRequest(_) + )); + } + + #[test] + fn unauthorized_rejection_remains_unauthorized() { + assert!(matches!( + public_error_for_rejection(InvocationRejectionReason::Unauthorized), + agent_error::Error::Unauthorized(_) + )); + } + + #[test] + fn internal_rejection_remains_internal() { + assert!(matches!( + public_error_for_rejection(InvocationRejectionReason::Internal), + agent_error::Error::InternalError(_) + )); + } + + #[derive(Clone)] + struct StaticShardManager(RoutingTable); + + #[async_trait::async_trait] + impl ShardManager for StaticShardManager { + async fn get_routing_table(&self) -> Result { + Ok(self.0.clone()) + } + + async fn register( + &self, + _port: u16, + _pod_name: Option, + ) -> Result { + unreachable!() + } + + async fn acquire_quota_lease( + &self, + _environment_id: EnvironmentId, + _resource_name: ResourceName, + _port: u16, + ) -> Result { + unreachable!() + } + + async fn renew_quota_lease( + &self, + _resource_definition_id: ResourceDefinitionId, + _port: u16, + _epoch: u64, + _unused: u64, + _pending_reservations: Vec, + ) -> Result { + unreachable!() + } + + async fn batch_renew_quota_leases( + &self, + _port: u16, + _renewals: Vec, + ) -> Result>, ShardManagerError> { + unreachable!() + } + + async fn release_quota_lease( + &self, + _resource_definition_id: ResourceDefinitionId, + _port: u16, + _epoch: u64, + _unused: u64, + ) -> Result<(), QuotaError> { + unreachable!() + } + } + + #[derive(Clone)] + struct RejectingExecutor; + + macro_rules! unimplemented_unary { + ($name:ident, $request:ty, $response:ty) => { + fn $name<'life0, 'async_trait>( + &'life0 self, + _request: Request<$request>, + ) -> Pin< + Box< + dyn std::future::Future, Status>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + }; + } + + #[tonic::async_trait] + impl WorkerExecutor for RejectingExecutor { + type ConnectWorkerStream = Pin< + Box< + dyn Stream> + + Send, + >, + >; + type GetFileContentsStream = + Pin> + Send>>; + type InvokeAgentSessionStream = + Pin> + Send>>; + + unimplemented_unary!(create_worker, CreateWorkerRequest, CreateWorkerResponse); + unimplemented_unary!(delete_worker, DeleteWorkerRequest, DeleteWorkerResponse); + unimplemented_unary!( + complete_promise, + CompletePromiseRequest, + CompletePromiseResponse + ); + unimplemented_unary!( + interrupt_worker, + InterruptWorkerRequest, + InterruptWorkerResponse + ); + unimplemented_unary!(revoke_shards, RevokeShardsRequest, RevokeShardsResponse); + unimplemented_unary!(assign_shards, AssignShardsRequest, AssignShardsResponse); + unimplemented_unary!( + set_shard_assignment, + SetShardAssignmentRequest, + SetShardAssignmentResponse + ); + unimplemented_unary!( + get_agent_metadata, + GetAgentMetadataRequest, + GetAgentMetadataResponse + ); + unimplemented_unary!(resume_worker, ResumeWorkerRequest, ResumeWorkerResponse); + unimplemented_unary!( + get_running_workers_metadata, + GetRunningWorkersMetadataRequest, + GetRunningWorkersMetadataResponse + ); + unimplemented_unary!( + get_workers_metadata, + GetWorkersMetadataRequest, + GetWorkersMetadataResponse + ); + unimplemented_unary!(update_worker, UpdateWorkerRequest, UpdateWorkerResponse); + unimplemented_unary!(get_oplog, GetOplogRequest, GetOplogResponse); + unimplemented_unary!(search_oplog, SearchOplogRequest, SearchOplogResponse); + unimplemented_unary!(fork_worker, ForkWorkerRequest, ForkWorkerResponse); + unimplemented_unary!(revert_worker, RevertWorkerRequest, RevertWorkerResponse); + unimplemented_unary!( + cancel_invocation, + CancelInvocationRequest, + CancelInvocationResponse + ); + unimplemented_unary!( + get_file_system_node, + GetFileSystemNodeRequest, + GetFileSystemNodeResponse + ); + unimplemented_unary!( + get_agent_wallet, + GetAgentWalletRequest, + GetAgentWalletResponse + ); + unimplemented_unary!( + activate_plugin, + ActivatePluginRequest, + ActivatePluginResponse + ); + unimplemented_unary!( + deactivate_plugin, + DeactivatePluginRequest, + DeactivatePluginResponse + ); + unimplemented_unary!( + process_oplog_entries, + ProcessOplogEntriesRequest, + ProcessOplogEntriesResponse + ); + + async fn connect_worker( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("connect_worker")) + } + + async fn get_file_contents( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("get_file_contents")) + } + + async fn invoke_agent_session( + &self, + request: Request>, + ) -> Result, Status> { + let mut requests = request.into_inner(); + let start = requests.message().await?.expect("missing invocation start"); + let (idempotency_key, agent_id) = match start.request { + Some(golem_api_grpc::proto::golem::worker::invocation_request::Request::Start( + start, + )) => (start.idempotency_key, start.agent_id), + other => panic!("expected invocation start, got {other:?}"), + }; + Ok(Response::new(Box::pin(stream::iter([Ok( + InvocationResponse { + response: Some(invocation_response::Response::Rejected( + InvocationRejected { + reason: InvocationRejectionReason::NotFound as i32, + error: "agent not found".to_string(), + idempotency_key, + agent_id, + component_revision: None, + }, + )), + }, + )])))) + } + } + + #[test] + async fn unary_not_found_rejection_preserves_the_public_error_category() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service( + WorkerExecutorServer::new(RejectingExecutor) + .accept_compressed(CompressionEncoding::Gzip) + .send_compressed(CompressionEncoding::Gzip), + ) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + + let routing_table: RoutingTable = GrpcRoutingTable { + number_of_shards: 1, + shard_assignments: vec![RoutingTableEntry { + shard_id: Some(ShardId { value: 0 }), + pod: Some(GrpcPod { + ip: Some(IpAddress { + kind: Some(ip_address::Kind::Ipv4(u32::from(Ipv4Addr::LOCALHOST))), + }), + port: port.into(), + }), + }], + } + .try_into() + .unwrap(); + let routing = Arc::new(RoutingTableService::new( + RoutingTableConfig::default(), + Arc::new(StaticShardManager(routing_table)), + )); + let clients = MultiTargetGrpcClient::new( + "provisional-rejecting-executor", + |channel: OtelGrpcService<_>, max_message_size| { + golem_api_grpc::proto::golem::workerexecutor::v1::worker_executor_client::WorkerExecutorClient::new(channel) + .send_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Gzip) + .max_decoding_message_size(max_message_size) + .max_encoding_message_size(max_message_size) + }, + GrpcClientConfig::default(), + ); + let client = WorkerExecutorWorkerClient::new(clients, RetryConfig::default(), routing); + let agent_id = AgentId { + component_id: ComponentId::new(), + agent_id: "missing".to_string(), + }; + + let error = client + .invoke_agent( + &agent_id, + Some("run".to_string()), + Some(SchemaValue { + value: Some(schema_value::Value::TupleValue(Default::default())), + }), + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + None, + Some(golem_common::model::IdempotencyKey::new( + "session-key".to_string(), + )), + None, + InvocationFreshnessDisposition::MayExist, + vec![], + EnvironmentId::new(), + AccountId::new(), + AuthCtx::System, + Principal::anonymous().into(), + ) + .await + .unwrap_err(); + + let public_error: AgentError = error.into(); + assert!( + matches!(public_error.error, Some(agent_error::Error::NotFound(_))), + "InvocationRejected(NotFound) must remain a public not-found error, got {public_error:?}" + ); + } +} diff --git a/golem-worker-service/src/service/worker/service.rs b/golem-worker-service/src/service/worker/service.rs index 3be7d2c6e4..cbab762aea 100644 --- a/golem-worker-service/src/service/worker/service.rs +++ b/golem-worker-service/src/service/worker/service.rs @@ -13,7 +13,10 @@ // limitations under the License. use super::WorkerResult; -use super::{ConnectWorkerStream, WorkerClient, WorkerServiceError}; +use super::{ + ConnectWorkerStream, InvocationRequestStream, InvocationResponseStream, WorkerClient, + WorkerServiceError, +}; use crate::api::agents::{ AgentInvocationMode, AgentInvocationRequest, AgentInvocationResult, CreateAgentRequest, CreateAgentResponse, @@ -23,14 +26,18 @@ use crate::service::auth::{AuthService, AuthServiceError}; use crate::service::component::ComponentService; use crate::service::limit::LimitService; use bytes::Bytes; -use futures::Stream; -use golem_api_grpc::proto::golem::worker::InvocationContext; +use futures::{Stream, StreamExt, stream}; +use golem_api_grpc::proto::golem::worker::{ + InvocationContext, InvocationRequest, InvocationStart, PublicInvocationStart, + invocation_request, +}; use golem_common::model::AgentInvocationOutput; use golem_common::model::account::AccountId; use golem_common::model::agent::{ AgentMode, AgentTypeName, GolemUserPrincipal, InvocationFreshnessDisposition, ParsedAgentId, Principal, ephemeral_invocation_phantom_id, }; +use golem_common::model::application::ApplicationName; use golem_common::model::card::owner::{AgentOwnerLeafPattern, AgentOwnerPattern}; use golem_common::model::card::{ AgentMethodName, AgentResourcePattern, AgentVerb, ClassPermissionTarget, PermissionTarget, @@ -40,7 +47,7 @@ use golem_common::model::component::{ CanonicalFilePath, ComponentId, ComponentName, ComponentRevision, PluginPriority, }; use golem_common::model::deployment::DeploymentRevision; -use golem_common::model::environment::EnvironmentId; +use golem_common::model::environment::{EnvironmentId, EnvironmentName}; use golem_common::model::oplog::OplogCursor; use golem_common::model::oplog::OplogIndex; use golem_common::model::worker::AgentConfigEntryDto; @@ -48,7 +55,12 @@ use golem_common::model::worker::AgentUpdateMode; use golem_common::model::worker::{AgentMetadataDto, RevertWorkerTarget}; use golem_common::model::{AgentFilter, AgentFingerprint, AgentId, IdempotencyKey, ScanCursor}; use golem_common::schema::json_input_schema_value_to_typed_schema_value; -use golem_common::schema::{SchemaType, TypedSchemaValue}; +use golem_common::schema::stream::SchemaValueStream; +use golem_common::schema::{ + ResultValuePayload, SchemaType, SchemaValue, TypedSchemaValue, UnionValuePayload, + VariantValuePayload, +}; +use golem_service_base::error::worker_executor::WorkerExecutorError; use golem_service_base::model::auth::AuthCtx; use golem_service_base::model::component::Component; use golem_service_base::model::{ComponentFileSystemNode, GetOplogResponse}; @@ -93,6 +105,188 @@ fn build_public_invocation_agent_id( }) } +fn validate_one_shot_invocation_is_stream_free( + component: &Component, + agent_id: &AgentId, + method_name: &str, + method_parameters: &golem_api_grpc::proto::golem::schema::SchemaValue, +) -> WorkerResult<()> { + let parsed_agent_id = ParsedAgentId::parse(&agent_id.agent_id, &component.metadata) + .map_err(WorkerServiceError::TypeChecker)?; + let agent_type = component + .metadata + .find_agent_type_by_name_ref(&parsed_agent_id.agent_type) + .ok_or_else(|| { + WorkerServiceError::TypeChecker(format!( + "Agent type '{}' not found", + parsed_agent_id.agent_type + )) + })?; + let method = agent_type + .methods + .iter() + .find(|method| method.name == method_name) + .ok_or_else(|| { + WorkerServiceError::TypeChecker(format!( + "Agent method '{method_name}' not found in agent type '{}'", + agent_type.type_name + )) + })?; + let input = SchemaValue::try_from(method_parameters.clone()) + .map_err(WorkerServiceError::TypeChecker)?; + method + .validate_input(&agent_type.schema, &input) + .map_err(|error| { + WorkerServiceError::TypeChecker(format!( + "Invalid input for agent method '{method_name}': {error}" + )) + })?; + if method.uses_streams(&agent_type.schema) { + Err(WorkerServiceError::TypeChecker( + "Streaming agent methods require an attached invocation session".to_string(), + )) + } else { + Ok(()) + } +} + +fn decode_public_session_schema_value( + value: golem_api_grpc::proto::golem::schema::SchemaValue, +) -> Result { + decode_public_schema_value(value, true) +} + +pub(crate) fn validate_public_session_schema_value( + value: &golem_api_grpc::proto::golem::schema::SchemaValue, +) -> Result<(), String> { + decode_public_session_schema_value(value.clone()).map(|_| ()) +} + +fn decode_public_schema_value( + value: golem_api_grpc::proto::golem::schema::SchemaValue, + allow_stream_references: bool, +) -> Result { + use golem_api_grpc::proto::golem::schema::{result_value, schema_value}; + + let value = value + .value + .ok_or_else(|| "Missing field: SchemaValue.value".to_string())?; + match value { + schema_value::Value::RecordValue(record) => Ok(SchemaValue::Record { + fields: record + .fields + .into_iter() + .map(|value| decode_public_schema_value(value, allow_stream_references)) + .collect::>()?, + }), + schema_value::Value::VariantValue(variant) => { + Ok(SchemaValue::Variant(VariantValuePayload { + case: variant.case, + payload: variant + .payload + .map(|payload| { + decode_public_schema_value(*payload, allow_stream_references).map(Box::new) + }) + .transpose()?, + })) + } + schema_value::Value::TupleValue(tuple) => Ok(SchemaValue::Tuple { + elements: tuple + .elements + .into_iter() + .map(|value| decode_public_schema_value(value, allow_stream_references)) + .collect::>()?, + }), + schema_value::Value::ListValue(list) => Ok(SchemaValue::List { + elements: list + .elements + .into_iter() + .map(|value| decode_public_schema_value(value, allow_stream_references)) + .collect::>()?, + }), + schema_value::Value::FixedListValue(list) => Ok(SchemaValue::FixedList { + elements: list + .elements + .into_iter() + .map(|value| decode_public_schema_value(value, allow_stream_references)) + .collect::>()?, + }), + schema_value::Value::MapValue(map) => Ok(SchemaValue::Map { + entries: map + .entries + .into_iter() + .map(|entry| { + let key = entry + .key + .ok_or_else(|| "Missing field: MapEntry.key".to_string())?; + let value = entry + .value + .ok_or_else(|| "Missing field: MapEntry.value".to_string())?; + Ok(( + decode_public_schema_value(key, allow_stream_references)?, + decode_public_schema_value(value, allow_stream_references)?, + )) + }) + .collect::>()?, + }), + schema_value::Value::OptionValue(option) => Ok(SchemaValue::Option { + inner: option + .inner + .map(|inner| { + decode_public_schema_value(*inner, allow_stream_references).map(Box::new) + }) + .transpose()?, + }), + schema_value::Value::ResultValue(result) => match result.result { + Some(result_value::Result::Ok(value)) => { + Ok(SchemaValue::Result(ResultValuePayload::Ok { + value: Some(Box::new(decode_public_schema_value( + *value, + allow_stream_references, + )?)), + })) + } + Some(result_value::Result::OkUnit(_)) => { + Ok(SchemaValue::Result(ResultValuePayload::Ok { value: None })) + } + Some(result_value::Result::Err(value)) => { + Ok(SchemaValue::Result(ResultValuePayload::Err { + value: Some(Box::new(decode_public_schema_value( + *value, + allow_stream_references, + )?)), + })) + } + Some(result_value::Result::ErrUnit(_)) => { + Ok(SchemaValue::Result(ResultValuePayload::Err { value: None })) + } + None => Err("Missing field: ResultValue.result".to_string()), + }, + schema_value::Value::UnionValue(union) => { + let body = union + .body + .ok_or_else(|| "Missing field: UnionValue.body".to_string())?; + Ok(SchemaValue::Union(UnionValuePayload { + tag: union.tag, + body: Box::new(decode_public_schema_value(*body, allow_stream_references)?), + })) + } + schema_value::Value::SecretValue(_) | schema_value::Value::QuotaTokenValue(_) => { + Err("host-managed capability values cannot cross the public boundary".to_string()) + } + schema_value::Value::StreamReference(reference) if allow_stream_references => Ok( + SchemaValue::Stream(SchemaValueStream::from_host_endpoint(reference.stream_id)), + ), + schema_value::Value::StreamReference(reference) => Err(format!( + "stream reference {} is not valid in constructor parameters", + reference.stream_id + )), + value => { + golem_api_grpc::proto::golem::schema::SchemaValue { value: Some(value) }.try_into() + } + } +} + fn normalize_agent_invocation_identity( component: &Component, agent_id: &AgentId, @@ -934,13 +1128,12 @@ impl WorkerService { config: Vec, auth_ctx: AuthCtx, principal: golem_api_grpc::proto::golem::component::Principal, - known_environment_id: Option, ) -> WorkerResult { let component = self .component_service .get_current_by_id(agent_id.component_id) .await?; - let environment_id = known_environment_id.unwrap_or(component.environment_id); + let environment_id = component.environment_id; let account_id = component.account_id; self.dispatch_agent_invocation( &component, @@ -976,6 +1169,283 @@ impl WorkerService { .await } + pub async fn invoke_agent_session( + &self, + mut start: InvocationStart, + tail: InvocationRequestStream, + allow_derived_ephemeral_phantom: bool, + auth_ctx: AuthCtx, + ) -> WorkerResult { + let agent_id: AgentId = start + .agent_id + .clone() + .ok_or_else(|| WorkerServiceError::TypeChecker("agent_id not found".to_string()))? + .try_into() + .map_err(WorkerServiceError::TypeChecker)?; + let component = self + .component_service + .get_current_by_id(agent_id.component_id) + .await?; + let observation_only = + start.mode() == golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup; + let freshness_disposition = if start.freshness_disposition() + == golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh + { + InvocationFreshnessDisposition::KnownFresh + } else { + InvocationFreshnessDisposition::MayExist + }; + let (agent_id, idempotency_key, mut freshness_disposition) = + normalize_agent_invocation_identity( + &component, + &agent_id, + start.idempotency_key.clone().map(Into::into), + allow_derived_ephemeral_phantom, + observation_only, + freshness_disposition, + )?; + if observation_only { + freshness_disposition = InvocationFreshnessDisposition::MayExist; + } + authorize_agent_permission( + &auth_ctx, + &component, + &agent_id, + agent_verb_for_invocation_mode(start.mode), + start + .method_name + .as_ref() + .map(|method_name| { + AgentResourcePattern::Method(AgentMethodName(method_name.clone())) + }) + .unwrap_or(AgentResourcePattern::Any), + )?; + + start.agent_id = Some(agent_id.clone().into()); + start.idempotency_key = Some(idempotency_key.into()); + start.environment_id = Some(component.environment_id.into()); + start.component_owner_account_id = Some(component.account_id.into()); + start.freshness_disposition = match freshness_disposition { + InvocationFreshnessDisposition::MayExist => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32 + } + InvocationFreshnessDisposition::KnownFresh => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh + as i32 + } + }; + let request = stream::once(async move { + InvocationRequest { + request: Some(invocation_request::Request::Start(start)), + } + }) + .chain(tail); + self.worker_client + .invoke_agent_session(&agent_id, Box::pin(request)) + .await + } + + pub async fn invoke_public_agent_session( + &self, + start: PublicInvocationStart, + tail: InvocationRequestStream, + auth: AuthCtx, + ) -> WorkerResult { + let app_name = ApplicationName::try_from(start.application_name) + .map_err(WorkerServiceError::TypeChecker)?; + let env_name = EnvironmentName::try_from(start.environment_name) + .map_err(WorkerServiceError::TypeChecker)?; + let agent_type_name = AgentTypeName(start.agent_type_name); + let method_name = start.method_name; + let constructor_parameters = decode_public_schema_value( + start.constructor_parameters.ok_or_else(|| { + WorkerServiceError::TypeChecker( + "public invocation has no constructor parameters".to_string(), + ) + })?, + false, + ) + .map_err(|error| { + WorkerServiceError::TypeChecker(format!( + "Agent constructor parameters cannot cross the public boundary: {error}" + )) + })?; + let proto_method_parameters = start.method_parameters.ok_or_else(|| { + WorkerServiceError::TypeChecker( + "public invocation has no method parameters".to_string(), + ) + })?; + let method_parameters = decode_public_session_schema_value(proto_method_parameters.clone()) + .map_err(|error| { + WorkerServiceError::TypeChecker(format!( + "Agent method parameters cannot cross the public boundary: {error}" + )) + })?; + let phantom_id = start + .phantom_id + .map(TryInto::try_into) + .transpose() + .map_err(|error| { + WorkerServiceError::TypeChecker(format!("Invalid phantom id: {error}")) + })?; + let idempotency_key: IdempotencyKey = start + .idempotency_key + .ok_or_else(|| { + WorkerServiceError::TypeChecker( + "public invocation requires an idempotency key".to_string(), + ) + })? + .into(); + let config = start + .config + .into_iter() + .map(AgentConfigEntryDto::try_from) + .collect::, _>>() + .map_err(|error| { + WorkerServiceError::TypeChecker(format!( + "Agent configuration cannot cross the public boundary: {error}" + )) + })?; + + let resolved = self + .agent_resolution_cache + .resolve(&app_name, &env_name, &agent_type_name, None, &auth) + .await?; + let registered_agent_type = &resolved.registered_agent_type; + let environment_id = resolved.environment_id; + let component_id = registered_agent_type.implemented_by.component_id; + let agent_type = ®istered_agent_type.agent_type; + + let constructor_parameters = json_input_schema_value_to_typed_schema_value( + constructor_parameters, + &agent_type.schema, + &agent_type.constructor.input_schema, + ) + .map_err(|error| { + WorkerServiceError::TypeChecker(format!( + "Agent constructor parameters type error: {error}" + )) + })?; + let agent_id = build_public_invocation_agent_id( + component_id, + agent_type_name.clone(), + constructor_parameters, + phantom_id, + )?; + let component = self + .component_service + .get_revision( + component_id, + registered_agent_type.implemented_by.component_revision, + ) + .await?; + let component_owner_account_id = registered_agent_type.implemented_by.account_id; + let component_name = registered_agent_type.implemented_by.component_name.clone(); + let component_owner_account_email = + registered_agent_type.implemented_by.account_email.clone(); + let (agent_id, idempotency_key, freshness_disposition, observation_only) = self + .prepare_agent_invocation_identity( + &component, + &agent_id, + Some(idempotency_key), + false, + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + InvocationFreshnessDisposition::MayExist, + |final_agent_id| { + auth.authorize_permission(&PermissionTarget::Agent(ClassPermissionTarget { + owner: AgentOwnerPattern::Agent { + account: component_owner_account_email, + application: app_name, + environment: env_name, + component: ComponentName(component_name), + agent: AgentOwnerLeafPattern::Agent(final_agent_id.agent_id.clone()), + }, + verb: Some(AgentVerb::Invoke), + resource: AgentResourcePattern::Method(AgentMethodName( + method_name.clone(), + )), + })) + .map_err(AuthServiceError::from) + .map_err(WorkerServiceError::from) + }, + )?; + debug_assert!(!observation_only); + let invocation_component = self + .component_for_invocation( + &component, + &agent_id, + environment_id, + &auth, + freshness_disposition, + ) + .await?; + let invocation_agent_type = invocation_component + .metadata + .find_agent_type_by_name_ref(&agent_type_name) + .ok_or_else(|| { + WorkerServiceError::Internal(format!( + "Agent type {agent_type_name} not found in component metadata at revision {}", + invocation_component.revision + )) + })?; + let method = invocation_agent_type + .methods + .iter() + .find(|method| method.name == method_name) + .ok_or_else(|| { + WorkerServiceError::Internal(format!( + "Agent method {method_name} not found in agent type {agent_type_name}" + )) + })?; + let _validated_method_parameters = json_input_schema_value_to_typed_schema_value( + method_parameters, + &invocation_agent_type.schema, + &method.input_schema, + ) + .map_err(|error| { + WorkerServiceError::TypeChecker(format!("Agent method parameters type error: {error}")) + })?; + let principal: golem_api_grpc::proto::golem::component::Principal = + Principal::GolemUser(GolemUserPrincipal { + account_id: auth.account_id(), + }) + .into(); + let trusted_start = InvocationStart { + agent_id: Some(agent_id.clone().into()), + method_name: Some(method_name), + input: Some(proto_method_parameters), + idempotency_key: Some(idempotency_key.into()), + context: None, + auth_ctx: Some(auth.into()), + principal: Some(principal), + environment_id: Some(environment_id.into()), + config: config.into_iter().map(Into::into).collect(), + component_owner_account_id: Some(component_owner_account_id.into()), + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + schedule_at: None, + freshness_disposition: match freshness_disposition { + InvocationFreshnessDisposition::MayExist => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32 + } + InvocationFreshnessDisposition::KnownFresh => { + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::KnownFresh + as i32 + } + }, + }; + let request = stream::once(async move { + InvocationRequest { + request: Some(invocation_request::Request::Start(trusted_start)), + } + }) + .chain(tail); + self.worker_client + .invoke_agent_session(&agent_id, Box::pin(request)) + .await + } + /// Shared invocation-dispatch core: normalizes the invocation identity, /// authorizes against the final agent id, dispatches to the executor, and /// backfills the final identity into the invocation output. @@ -999,6 +1469,67 @@ impl WorkerService { principal: golem_api_grpc::proto::golem::component::Principal, authorize: impl FnOnce(&AgentId) -> WorkerResult<()>, ) -> WorkerResult { + let (agent_id, idempotency_key, freshness_disposition, observation_only) = self + .prepare_agent_invocation_identity( + component, + agent_id, + idempotency_key, + allow_derived_ephemeral_phantom, + mode, + freshness_disposition, + authorize, + )?; + + let validation_component = if observation_only { + None + } else { + Some( + self.component_for_invocation( + component, + &agent_id, + environment_id, + &auth_ctx, + freshness_disposition, + ) + .await?, + ) + }; + + self.dispatch_prepared_agent_invocation( + validation_component.as_ref(), + agent_id, + method_name, + method_parameters, + mode, + schedule_at, + idempotency_key, + invocation_context, + freshness_disposition, + config, + environment_id, + account_id, + auth_ctx, + principal, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + fn prepare_agent_invocation_identity( + &self, + component: &Component, + agent_id: &AgentId, + idempotency_key: Option, + allow_derived_ephemeral_phantom: bool, + mode: i32, + freshness_disposition: InvocationFreshnessDisposition, + authorize: impl FnOnce(&AgentId) -> WorkerResult<()>, + ) -> WorkerResult<( + AgentId, + IdempotencyKey, + InvocationFreshnessDisposition, + bool, + )> { let observation_only = mode == golem_api_grpc::proto::golem::worker::AgentInvocationMode::Lookup as i32; let (agent_id, idempotency_key, mut freshness_disposition) = @@ -1014,6 +1545,50 @@ impl WorkerService { freshness_disposition = InvocationFreshnessDisposition::MayExist; } authorize(&agent_id)?; + Ok(( + agent_id, + idempotency_key, + freshness_disposition, + observation_only, + )) + } + + #[allow(clippy::too_many_arguments)] + async fn dispatch_prepared_agent_invocation( + &self, + validation_component: Option<&Component>, + agent_id: AgentId, + method_name: Option, + method_parameters: Option, + mode: i32, + schedule_at: Option<::prost_types::Timestamp>, + idempotency_key: IdempotencyKey, + invocation_context: Option, + freshness_disposition: InvocationFreshnessDisposition, + config: Vec, + environment_id: EnvironmentId, + account_id: AccountId, + auth_ctx: AuthCtx, + principal: golem_api_grpc::proto::golem::component::Principal, + ) -> WorkerResult { + if let Some(validation_component) = validation_component { + let method_name = method_name.as_deref().ok_or_else(|| { + WorkerServiceError::TypeChecker( + "method_name is required for non-lookup invocations".to_string(), + ) + })?; + let method_parameters = method_parameters.as_ref().ok_or_else(|| { + WorkerServiceError::TypeChecker( + "method_parameters are required for non-lookup invocations".to_string(), + ) + })?; + validate_one_shot_invocation_is_stream_free( + validation_component, + &agent_id, + method_name, + method_parameters, + )?; + } let mut output = self .worker_client @@ -1038,6 +1613,41 @@ impl WorkerService { Ok(output) } + async fn component_for_invocation( + &self, + fallback: &Component, + agent_id: &AgentId, + environment_id: EnvironmentId, + auth_ctx: &AuthCtx, + freshness_disposition: InvocationFreshnessDisposition, + ) -> WorkerResult { + if freshness_disposition == InvocationFreshnessDisposition::KnownFresh { + return Ok(fallback.clone()); + } + + let component_revision = match self + .worker_client + .get_metadata(agent_id, environment_id, auth_ctx.clone()) + .await + { + Ok(metadata) => metadata.component_revision, + Err(WorkerServiceError::AgentNotFound(_)) + | Err(WorkerServiceError::GolemError(WorkerExecutorError::AgentNotFound { .. })) => { + return Ok(fallback.clone()); + } + Err(error) => return Err(error), + }; + + if component_revision == fallback.revision { + Ok(fallback.clone()) + } else { + Ok(self + .component_service + .get_revision(fallback.id, component_revision) + .await?) + } + } + /// REST path: resolves the agent via the registry, validates its parameters, then creates it. pub async fn create_agent_rest( &self, @@ -1181,13 +1791,67 @@ impl WorkerService { registered_agent_type.implemented_by.component_revision, ) .await?; + let component_name = registered_agent_type.implemented_by.component_name.clone(); + let component_owner_account_id = registered_agent_type.implemented_by.account_id; + let component_owner_account_email = + registered_agent_type.implemented_by.account_email.clone(); + let method_name = request.method_name.clone(); + let agent_type_name = request.agent_type_name.clone(); + let proto_mode = match request.mode { + AgentInvocationMode::Await => { + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32 + } + AgentInvocationMode::Schedule => { + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32 + } + }; + let (agent_id, idempotency_key, freshness_disposition, observation_only) = self + .prepare_agent_invocation_identity( + &component, + &agent_id, + request.idempotency_key.clone(), + false, + proto_mode, + InvocationFreshnessDisposition::MayExist, + |final_agent_id| { + auth.authorize_permission(&PermissionTarget::Agent(ClassPermissionTarget { + owner: AgentOwnerPattern::Agent { + account: component_owner_account_email, + application: request.app_name, + environment: request.env_name, + component: ComponentName(component_name), + agent: AgentOwnerLeafPattern::Agent(final_agent_id.agent_id.clone()), + }, + verb: Some(AgentVerb::Invoke), + resource: AgentResourcePattern::Method(AgentMethodName( + method_name.clone(), + )), + })) + .map_err(AuthServiceError::from) + .map_err(WorkerServiceError::from) + }, + )?; + debug_assert!(!observation_only); + let invocation_component = self + .component_for_invocation( + &component, + &agent_id, + environment_id, + &auth, + freshness_disposition, + ) + .await?; - let component_name = registered_agent_type.implemented_by.component_name.clone(); - let component_owner_account_id = registered_agent_type.implemented_by.account_id; - let component_owner_account_email = - registered_agent_type.implemented_by.account_email.clone(); - - let method = agent_type + let invocation_agent_type = invocation_component + .metadata + .find_agent_type_by_name_ref(&request.agent_type_name) + .ok_or_else(|| { + WorkerServiceError::Internal(format!( + "Agent type {} not found in component metadata at revision {}", + request.agent_type_name, invocation_component.revision + )) + })?; + let method = invocation_agent_type .methods .iter() .find(|m| m.name == request.method_name) @@ -1200,7 +1864,7 @@ impl WorkerService { let method_parameters = json_input_schema_value_to_typed_schema_value( request.method_parameters, - &agent_type.schema, + &invocation_agent_type.schema, &method.input_schema, ) .map_err(|err| { @@ -1210,16 +1874,11 @@ impl WorkerService { .1; let proto_method_parameters: golem_api_grpc::proto::golem::schema::SchemaValue = - method_parameters.into(); - - let proto_mode = match request.mode { - AgentInvocationMode::Await => { - golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32 - } - AgentInvocationMode::Schedule => { - golem_api_grpc::proto::golem::worker::AgentInvocationMode::Schedule as i32 - } - }; + method_parameters.try_into().map_err(|error| { + WorkerServiceError::TypeChecker(format!( + "Agent method parameters cannot cross the worker boundary: {error}" + )) + })?; let proto_schedule_at = request.schedule_at.map(|dt| ::prost_types::Timestamp { seconds: dt.timestamp(), @@ -1232,43 +1891,22 @@ impl WorkerService { }) .into(); - let method_name = request.method_name.clone(); - let agent_type_name = request.agent_type_name.clone(); - let output = self - .dispatch_agent_invocation( - &component, - &agent_id, + .dispatch_prepared_agent_invocation( + Some(&invocation_component), + agent_id.clone(), Some(method_name.clone()), Some(proto_method_parameters), proto_mode, proto_schedule_at, - request.idempotency_key.clone(), + idempotency_key, None, - false, - InvocationFreshnessDisposition::MayExist, + freshness_disposition, request.config, environment_id, component_owner_account_id, auth.clone(), principal, - |final_agent_id| { - auth.authorize_permission(&PermissionTarget::Agent(ClassPermissionTarget { - owner: AgentOwnerPattern::Agent { - account: component_owner_account_email, - application: request.app_name, - environment: request.env_name, - component: ComponentName(component_name), - agent: AgentOwnerLeafPattern::Agent(final_agent_id.agent_id.clone()), - }, - verb: Some(AgentVerb::Invoke), - resource: AgentResourcePattern::Method(AgentMethodName( - method_name.clone(), - )), - })) - .map_err(AuthServiceError::from) - .map_err(WorkerServiceError::from) - }, ) .await?; @@ -1334,7 +1972,8 @@ impl WorkerService { mod tests { use super::{ WorkerService, agent_verb_for_invocation_mode, build_public_agent_id, - build_public_invocation_agent_id, normalize_agent_invocation_identity, + build_public_invocation_agent_id, decode_public_schema_value, + normalize_agent_invocation_identity, }; use crate::api::agents::{AgentInvocationMode, AgentInvocationRequest, CreateAgentRequest}; use crate::service::agent_resolution_cache::AgentResolutionCache; @@ -1345,8 +1984,10 @@ mod tests { use async_trait::async_trait; use bytes::Bytes; use chrono::Utc; - use futures::Stream; - use golem_api_grpc::proto::golem::worker::{InvocationContext, LogEvent}; + use futures::{Stream, StreamExt, stream}; + use golem_api_grpc::proto::golem::worker::{ + InvocationContext, InvocationStart, LogEvent, PublicInvocationStart, invocation_request, + }; use golem_common::base_model::component_metadata::KnownExports; use golem_common::model::AgentInvocationOutput; use golem_common::model::Empty; @@ -1357,7 +1998,7 @@ mod tests { ResolvedAgentType, Snapshotting, ephemeral_invocation_phantom_id, }; use golem_common::model::application::{ApplicationId, ApplicationName}; - use golem_common::model::card::{AgentVerb, StoredCard}; + use golem_common::model::card::{AgentVerb, EffectiveSurface, StoredCard}; use golem_common::model::component::{ CanonicalFilePath, ComponentId, ComponentName, ComponentRevision, PluginPriority, }; @@ -1367,16 +2008,18 @@ mod tests { use golem_common::model::environment::{EnvironmentId, EnvironmentName}; use golem_common::model::oplog::{OplogCursor, OplogIndex}; use golem_common::model::worker::{AgentConfigEntryDto, AgentMetadataDto, RevertWorkerTarget}; - use golem_common::model::{AgentFilter, AgentFingerprint, AgentId, IdempotencyKey, ScanCursor}; + use golem_common::model::{ + AgentFilter, AgentFingerprint, AgentId, AgentStatus, IdempotencyKey, ScanCursor, Timestamp, + }; use golem_common::schema::{ - AgentConstructorSchema, AgentMethodSchema, AgentTypeSchema, InputSchema, OutputSchema, - SchemaGraph, SchemaValue, + AgentConstructorSchema, AgentMethodSchema, AgentTypeSchema, InputSchema, NamedField, + OutputSchema, SchemaGraph, SchemaType, SchemaValue, }; use golem_service_base::clients::registry::{RegistryService, RegistryServiceError}; use golem_service_base::model::auth::AuthCtx; use golem_service_base::model::component::Component; use golem_service_base::model::{ComponentFileSystemNode, GetOplogResponse}; - use std::collections::{BTreeMap, HashMap}; + use std::collections::{BTreeMap, HashMap, HashSet}; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -1770,24 +2413,26 @@ mod tests { } struct StaticComponentService { - component: Component, + components: Vec, } #[async_trait] impl ComponentService for StaticComponentService { async fn get_current_by_id_in_cache(&self, component_id: ComponentId) -> Option { - (self.component.id == component_id).then(|| self.component.clone()) + self.components + .iter() + .filter(|component| component.id == component_id) + .max_by_key(|component| component.revision) + .cloned() } async fn get_current_by_id_uncached( &self, component_id: ComponentId, ) -> Result { - if self.component.id == component_id { - Ok(self.component.clone()) - } else { - Err(ComponentServiceError::ComponentNotFound) - } + self.get_current_by_id_in_cache(component_id) + .await + .ok_or(ComponentServiceError::ComponentNotFound) } async fn get_revision( @@ -1795,21 +2440,29 @@ mod tests { component_id: ComponentId, component_revision: ComponentRevision, ) -> Result { - if self.component.id == component_id && self.component.revision == component_revision { - Ok(self.component.clone()) - } else { - Err(ComponentServiceError::ComponentNotFound) - } + self.components + .iter() + .find(|component| { + component.id == component_id && component.revision == component_revision + }) + .cloned() + .ok_or(ComponentServiceError::ComponentNotFound) } async fn get_all_revisions( &self, component_id: ComponentId, ) -> Result, ComponentServiceError> { - if self.component.id == component_id { - Ok(vec![self.component.clone()]) - } else { + let components = self + .components + .iter() + .filter(|component| component.id == component_id) + .cloned() + .collect::>(); + if components.is_empty() { Err(ComponentServiceError::ComponentNotFound) + } else { + Ok(components) } } } @@ -1843,7 +2496,10 @@ mod tests { struct RecordingWorkerClient { created_agent_ids: Mutex>, invocations: Mutex>, + invocation_environments: Mutex>, + invocation_session_starts: Mutex>, invocation_output: AgentInvocationOutput, + metadata_component_revision: Option, } impl RecordingWorkerClient { @@ -1851,7 +2507,24 @@ mod tests { Self { created_agent_ids: Mutex::new(Vec::new()), invocations: Mutex::new(Vec::new()), + invocation_environments: Mutex::new(Vec::new()), + invocation_session_starts: Mutex::new(Vec::new()), + invocation_output, + metadata_component_revision: None, + } + } + + fn with_metadata_component_revision( + invocation_output: AgentInvocationOutput, + component_revision: ComponentRevision, + ) -> Self { + Self { + created_agent_ids: Mutex::new(Vec::new()), + invocations: Mutex::new(Vec::new()), + invocation_environments: Mutex::new(Vec::new()), + invocation_session_starts: Mutex::new(Vec::new()), invocation_output, + metadata_component_revision: Some(component_revision), } } @@ -1866,6 +2539,18 @@ mod tests { fn invocations(&self) -> Vec<(AgentId, IdempotencyKey, InvocationFreshnessDisposition)> { self.invocations.lock().unwrap().clone() } + + fn invocation_environment(&self) -> EnvironmentId { + self.invocation_environments.lock().unwrap()[0] + } + + fn invocation_session_start(&self) -> (AgentId, InvocationStart) { + self.invocation_session_starts.lock().unwrap()[0].clone() + } + + fn invocation_session_start_count(&self) -> usize { + self.invocation_session_starts.lock().unwrap().len() + } } #[async_trait] @@ -1926,11 +2611,35 @@ mod tests { async fn get_metadata( &self, - _: &AgentId, - _: EnvironmentId, + agent_id: &AgentId, + environment_id: EnvironmentId, _: AuthCtx, ) -> WorkerResult { - unimplemented!() + match self.metadata_component_revision { + Some(component_revision) => Ok(AgentMetadataDto { + agent_id: agent_id.clone(), + environment_id, + created_by: AccountId(Uuid::new_v4()), + env: HashMap::new(), + config: Vec::new(), + status: AgentStatus::Idle, + component_revision, + retry_count: 0, + pending_invocation_count: 0, + updates: Vec::new(), + created_at: Timestamp::now_utc(), + last_error: None, + component_size: 0, + total_linear_memory_size: 0, + exported_resource_instances: Vec::new(), + active_plugins: HashSet::new(), + skipped_regions: Vec::new(), + deleted_regions: Vec::new(), + last_oplog_index: OplogIndex::INITIAL, + fingerprint: AgentFingerprint(Uuid::new_v4()), + }), + None => Err(WorkerServiceError::AgentNotFound(agent_id.clone())), + } } async fn find_metadata( @@ -2089,7 +2798,7 @@ mod tests { _: Option, freshness_disposition: InvocationFreshnessDisposition, _: Vec, - _: EnvironmentId, + environment_id: EnvironmentId, _: AccountId, _: AuthCtx, _: golem_api_grpc::proto::golem::component::Principal, @@ -2099,9 +2808,33 @@ mod tests { idempotency_key.expect("worker service should supply an idempotency key"), freshness_disposition, )); + self.invocation_environments + .lock() + .unwrap() + .push(environment_id); Ok(self.invocation_output.clone()) } + async fn invoke_agent_session( + &self, + agent_id: &AgentId, + mut request: super::InvocationRequestStream, + ) -> WorkerResult { + let start = request + .next() + .await + .expect("worker service should send an invocation start"); + let start = match start.request { + Some(invocation_request::Request::Start(start)) => start, + other => panic!("expected invocation start, got {other:?}"), + }; + self.invocation_session_starts + .lock() + .unwrap() + .push((agent_id.clone(), start)); + Ok(Box::pin(stream::empty())) + } + async fn process_oplog_entries( &self, _: &AgentId, @@ -2125,16 +2858,35 @@ mod tests { agent_type_name: AgentTypeName, component_id: ComponentId, component_revision: ComponentRevision, + environment_id: EnvironmentId, } impl RestHarness { fn new(mode: AgentMode) -> Self { + Self::new_with_method_schema(mode, InputSchema::Parameters(vec![]), OutputSchema::Unit) + } + + fn new_with_output(mode: AgentMode, output_schema: OutputSchema) -> Self { + Self::new_with_method_schema(mode, InputSchema::Parameters(vec![]), output_schema) + } + + fn new_with_input(mode: AgentMode, input_schema: InputSchema) -> Self { + Self::new_with_method_schema(mode, input_schema, OutputSchema::Unit) + } + + fn new_with_method_schema( + mode: AgentMode, + input_schema: InputSchema, + output_schema: OutputSchema, + ) -> Self { let component_id = ComponentId(Uuid::new_v4()); let environment_id = EnvironmentId(Uuid::new_v4()); let account_id = AccountId(Uuid::new_v4()); let component_revision = ComponentRevision::INITIAL; let agent_type_name = AgentTypeName("weather-agent".to_string()); - let agent_type = test_agent_type(agent_type_name.clone(), mode); + let mut agent_type = test_agent_type(agent_type_name.clone(), mode); + agent_type.methods[0].input_schema = input_schema; + agent_type.methods[0].output_schema = output_schema; let component = test_component( component_id, environment_id, @@ -2178,7 +2930,9 @@ mod tests { Self { worker_service: WorkerService::new( - Arc::new(StaticComponentService { component }), + Arc::new(StaticComponentService { + components: vec![component], + }), Arc::new(AllowAllAuthService), Arc::new(NoopLimitService), worker_client.clone(), @@ -2188,6 +2942,92 @@ mod tests { agent_type_name, component_id, component_revision, + environment_id, + } + } + + fn new_with_pinned_and_latest_output( + pinned_output_schema: OutputSchema, + latest_output_schema: OutputSchema, + ) -> Self { + let component_id = ComponentId(Uuid::new_v4()); + let environment_id = EnvironmentId(Uuid::new_v4()); + let account_id = AccountId(Uuid::new_v4()); + let pinned_revision = ComponentRevision::INITIAL; + let latest_revision = ComponentRevision::new(1).unwrap(); + let agent_type_name = AgentTypeName("weather-agent".to_string()); + let mut pinned_agent_type = + test_agent_type(agent_type_name.clone(), AgentMode::Durable); + pinned_agent_type.methods[0].output_schema = pinned_output_schema; + let mut latest_agent_type = + test_agent_type(agent_type_name.clone(), AgentMode::Durable); + latest_agent_type.methods[0].output_schema = latest_output_schema; + let pinned_component = test_component( + component_id, + environment_id, + account_id, + pinned_revision, + pinned_agent_type, + ); + let latest_component = test_component( + component_id, + environment_id, + account_id, + latest_revision, + latest_agent_type.clone(), + ); + let worker_client = Arc::new(RecordingWorkerClient::with_metadata_component_revision( + AgentInvocationOutput { + result: golem_common::model::AgentInvocationResult::AgentInitialization, + consumed_fuel: None, + invocation_status: None, + component_revision: Some(pinned_revision), + agent_id: None, + idempotency_key: None, + oplog_index: None, + agent_fingerprint: None, + }, + pinned_revision, + )); + let registry = Arc::new(TestRegistryService { + resolved: ResolvedAgentType { + registered_agent_type: RegisteredAgentType { + agent_type: latest_agent_type, + implemented_by: RegisteredAgentTypeImplementer { + component_id, + component_revision: latest_revision, + component_name: latest_component.component_name.0.clone(), + account_id: latest_component.account_id, + account_email: latest_component.account_email.clone(), + }, + }, + environment_id, + deployment_revision: DeploymentRevision::INITIAL, + current_deployment_revision: Some(CurrentDeploymentRevision::INITIAL), + }, + }); + let agent_resolution_cache = Arc::new(AgentResolutionCache::new( + registry, + 1, + Duration::from_secs(60), + Duration::from_secs(60), + )); + + Self { + worker_service: WorkerService::new( + Arc::new(StaticComponentService { + components: vec![pinned_component, latest_component], + }), + Arc::new(AllowAllAuthService), + Arc::new(NoopLimitService), + worker_client.clone(), + agent_resolution_cache, + ), + worker_client, + agent_type_name, + component_id, + component_revision: latest_revision, + environment_id, } } @@ -2219,6 +3059,23 @@ mod tests { owner_account_email: None, } } + + fn public_invocation_start( + &self, + idempotency_key: IdempotencyKey, + ) -> PublicInvocationStart { + PublicInvocationStart { + application_name: "weather-app".to_string(), + environment_name: "prod".to_string(), + agent_type_name: self.agent_type_name.0.clone(), + constructor_parameters: Some(empty_json_tuple().try_into().unwrap()), + phantom_id: None, + config: vec![], + method_name: "run".to_string(), + method_parameters: Some(empty_json_tuple().try_into().unwrap()), + idempotency_key: Some(idempotency_key.into()), + } + } } fn test_agent_type(agent_type_name: AgentTypeName, mode: AgentMode) -> AgentTypeSchema { @@ -2356,6 +3213,10 @@ mod tests { Some(harness.component_revision) ); assert_eq!(response.agent_id, harness.worker_client.invoked_agent_id()); + assert_eq!( + harness.worker_client.invocation_environment(), + harness.environment_id + ); assert!(phantom_id(&response.agent_id).is_some()); let invocations = harness.worker_client.invocations(); assert_eq!(invocations.len(), 1); @@ -2367,6 +3228,236 @@ mod tests { ); } + #[test] + async fn public_invocation_session_resolves_normalizes_and_builds_trusted_start() { + let harness = RestHarness::new(AgentMode::Ephemeral); + let idempotency_key = IdempotencyKey::new("public-session-key".to_string()); + + let _responses = harness + .worker_service + .invoke_public_agent_session( + harness.public_invocation_start(idempotency_key.clone()), + Box::pin(stream::empty()), + AuthCtx::system(), + ) + .await + .unwrap(); + + let (routed_agent_id, start) = harness.worker_client.invocation_session_start(); + let trusted_agent_id: AgentId = start.agent_id.clone().unwrap().try_into().unwrap(); + let trusted_idempotency_key: IdempotencyKey = start.idempotency_key.clone().unwrap().into(); + + assert_eq!(routed_agent_id, trusted_agent_id); + assert_eq!(trusted_agent_id.component_id, harness.component_id); + assert_eq!(trusted_idempotency_key, idempotency_key); + assert_eq!( + phantom_id(&trusted_agent_id), + Some(ephemeral_invocation_phantom_id(&idempotency_key)) + ); + assert_eq!(start.method_name.as_deref(), Some("run")); + assert!(start.input.is_some()); + assert!(start.auth_ctx.is_some()); + assert!(start.principal.is_some()); + assert_eq!( + EnvironmentId::try_from(start.environment_id.unwrap()).unwrap(), + harness.environment_id + ); + assert!(start.component_owner_account_id.is_some()); + assert_eq!( + start.mode(), + golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await + ); + assert!(start.schedule_at.is_none()); + assert_eq!( + start.freshness_disposition(), + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + ); + } + + #[test] + async fn public_invocation_session_validates_and_preserves_live_stream_references() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue as ProtoSchemaValue, SchemaValueStreamReference, schema_value, + }; + + let harness = RestHarness::new_with_input( + AgentMode::Durable, + InputSchema::Parameters(vec![NamedField::user_supplied( + "input", + SchemaType::stream(Some(SchemaType::u32())), + )]), + ); + let method_parameters = ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 1 }, + )), + }], + })), + }; + let mut start = harness.public_invocation_start(IdempotencyKey::fresh()); + start.method_parameters = Some(method_parameters.clone()); + + let _responses = harness + .worker_service + .invoke_public_agent_session(start, Box::pin(stream::empty()), AuthCtx::system()) + .await + .unwrap(); + + let (_, trusted_start) = harness.worker_client.invocation_session_start(); + assert_eq!(trusted_start.input, Some(method_parameters)); + } + + #[test] + fn public_invocation_values_reject_capabilities_and_constructor_streams_recursively() { + use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue as ProtoSchemaValue, SchemaValueStreamReference, SecretValue, + schema_value, + }; + + let nested = |value| ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { value: Some(value) }], + })), + }; + + let secret = nested(schema_value::Value::SecretValue(SecretValue::default())); + assert!( + decode_public_schema_value(secret, true) + .unwrap_err() + .contains("host-managed capability") + ); + + let constructor_stream = nested(schema_value::Value::StreamReference( + SchemaValueStreamReference { stream_id: 7 }, + )); + assert!( + decode_public_schema_value(constructor_stream, false) + .unwrap_err() + .contains("not valid in constructor parameters") + ); + } + + #[test] + async fn unauthorized_public_invocation_session_never_reaches_worker_dispatch() { + let harness = RestHarness::new(AgentMode::Durable); + let idempotency_key = IdempotencyKey::fresh(); + let auth = AuthCtx::agent_with_effective_surface( + AccountId(Uuid::new_v4()), + AccountEmail::new("unauthorized@golem"), + EffectiveSurface { + source_card_ids: vec![], + lower: vec![], + upper: vec![], + }, + ); + + let error = match harness + .worker_service + .invoke_public_agent_session( + harness.public_invocation_start(idempotency_key), + Box::pin(stream::empty()), + auth, + ) + .await + { + Err(error) => error, + Ok(_) => panic!("an empty permission surface must reject the resolved agent selector"), + }; + + assert!(matches!(error, WorkerServiceError::AuthError(_))); + assert_eq!(harness.worker_client.invocation_session_start_count(), 0); + } + + #[test] + async fn non_attached_streaming_modes_are_rejected_before_worker_dispatch() { + let harness = RestHarness::new_with_output( + AgentMode::Durable, + OutputSchema::Single(Box::new(SchemaType::stream(Some(SchemaType::u8())))), + ); + + for (mode, schedule_at) in [ + (AgentInvocationMode::Await, None), + (AgentInvocationMode::Schedule, None), + (AgentInvocationMode::Schedule, Some(Utc::now())), + ] { + let mut request = harness.invoke_request(); + request.mode = mode; + request.schedule_at = schedule_at; + let error = harness + .worker_service + .invoke_agent_rest(request, AuthCtx::system()) + .await + .expect_err("non-attached invocation must reject streaming methods"); + + assert!( + error + .to_string() + .contains("require an attached invocation session"), + "unexpected error: {error}" + ); + } + assert!( + harness.worker_client.invocations().is_empty(), + "rejection must happen before scheduling, executor dispatch, enqueue, or result storage" + ); + } + + #[test] + async fn non_attached_stream_free_modes_still_dispatch() { + let harness = RestHarness::new(AgentMode::Durable); + + for (mode, schedule_at) in [ + (AgentInvocationMode::Await, None), + (AgentInvocationMode::Schedule, None), + (AgentInvocationMode::Schedule, Some(Utc::now())), + ] { + let mut request = harness.invoke_request(); + request.mode = mode; + request.schedule_at = schedule_at; + harness + .worker_service + .invoke_agent_rest(request, AuthCtx::system()) + .await + .expect("stream-free one-shot invocation must still dispatch"); + } + + assert_eq!(harness.worker_client.invocations().len(), 3); + } + + #[test] + async fn non_attached_classification_uses_the_existing_workers_component_revision() { + let stream_output = + OutputSchema::Single(Box::new(SchemaType::stream(Some(SchemaType::u8())))); + let pinned_streaming = RestHarness::new_with_pinned_and_latest_output( + stream_output.clone(), + OutputSchema::Unit, + ); + + let error = pinned_streaming + .worker_service + .invoke_agent_rest(pinned_streaming.invoke_request(), AuthCtx::system()) + .await + .expect_err("the pinned streaming revision must be rejected"); + assert!( + error + .to_string() + .contains("require an attached invocation session"), + "unexpected error: {error}" + ); + assert!(pinned_streaming.worker_client.invocations().is_empty()); + + let pinned_stream_free = + RestHarness::new_with_pinned_and_latest_output(OutputSchema::Unit, stream_output); + pinned_stream_free + .worker_service + .invoke_agent_rest(pinned_stream_free.invoke_request(), AuthCtx::system()) + .await + .expect("the pinned stream-free revision must remain dispatchable"); + assert_eq!(pinned_stream_free.worker_client.invocations().len(), 1); + } + #[test] async fn ephemeral_rest_invocations_get_fresh_identities_per_request() { let harness = RestHarness::new(AgentMode::Ephemeral); @@ -2475,7 +3566,6 @@ mod tests { Vec::new(), AuthCtx::system(), Principal::anonymous().into(), - None, ) .await .unwrap(); diff --git a/integration-tests/Cargo.toml b/integration-tests/Cargo.toml index acc64ba4c7..eaebd721c9 100644 --- a/integration-tests/Cargo.toml +++ b/integration-tests/Cargo.toml @@ -31,6 +31,7 @@ futures-concurrency = { workspace = true } headers = { workspace = true } indoc = { workspace = true } pretty_assertions = { workspace = true, features = [ "unstable" ] } +prost = { workspace = true } rand = { workspace = true } reqwest = { workspace = true } rlimit = { workspace = true } diff --git a/integration-tests/tests/api/mod.rs b/integration-tests/tests/api/mod.rs index 9e1e068c90..df144c20df 100644 --- a/integration-tests/tests/api/mod.rs +++ b/integration-tests/tests/api/mod.rs @@ -30,6 +30,7 @@ mod resource_definition; mod retry_policies; mod rpc_auth; mod security_schemes; +mod streaming_rpc; use golem_test_framework::config::EnvBasedTestDependencies; use test_r::inherit_test_dep; diff --git a/integration-tests/tests/api/streaming_rpc.rs b/integration-tests/tests/api/streaming_rpc.rs new file mode 100644 index 0000000000..07d6939cd6 --- /dev/null +++ b/integration-tests/tests/api/streaming_rpc.rs @@ -0,0 +1,807 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use futures::{SinkExt, StreamExt}; +use golem_api_grpc::invocation_session_protocol::InvocationSessionState; +use golem_api_grpc::proto::golem::schema::{ + RecordValue, SchemaValue as ProtoSchemaValue, SchemaValueStreamReference, SecretValue, + schema_value, +}; +use golem_api_grpc::proto::golem::worker::v1::worker_service_client::WorkerServiceClient; +use golem_api_grpc::proto::golem::worker::{ + InputStreamEnd, InputStreamItem, InvocationRejectionReason, InvocationRequest, + InvocationResponse, InvocationStart, PublicInvocationRequest, PublicInvocationStart, + input_stream_item, invocation_request, invocation_response, invocation_session_completion, + invocation_session_result, public_invocation_request, +}; +use golem_client::model::ComponentDto; +use golem_common::model::agent::ParsedAgentId; +use golem_common::model::auth::TokenSecret; +use golem_common::model::{AgentId, IdempotencyKey, RoutingTable}; +use golem_common::schema::{SchemaValue, TypedSchemaValue}; +use golem_common::{agent_id, data_value}; +use golem_service_base::model::auth::AuthCtx; +use golem_test_framework::config::{EnvBasedTestDependencies, TestDependencies}; +use golem_test_framework::dsl::{TestDsl, TestDslExtended}; +use prost::Message as ProstMessage; +use test_r::{inherit_test_dep, test, timeout}; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::tungstenite::{Error as WebSocketError, Message}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +inherit_test_dep!(EnvBasedTestDependencies); + +type PublicInvocationSocket = WebSocketStream>; + +async fn connect_public_invocation_socket( + deps: &EnvBasedTestDependencies, + token: Option<&TokenSecret>, +) -> Result { + let worker_service = deps.worker_service(); + let url = format!( + "ws://{}:{}/v1/agents/invoke-agent-session", + worker_service.http_host(), + worker_service.http_port() + ); + let mut request = url.into_client_request()?; + if let Some(token) = token { + request.headers_mut().insert( + AUTHORIZATION, + format!("Bearer {}", token.secret()).parse().unwrap(), + ); + } + tokio_tungstenite::connect_async(request) + .await + .map(|(socket, _)| socket) +} + +async fn send_public_request( + socket: &mut PublicInvocationSocket, + request: &PublicInvocationRequest, +) -> anyhow::Result<()> { + socket + .send(Message::Binary(request.encode_to_vec().into())) + .await?; + Ok(()) +} + +async fn receive_public_response( + socket: &mut PublicInvocationSocket, +) -> anyhow::Result { + loop { + match socket.next().await { + Some(Ok(Message::Binary(payload))) => { + return InvocationResponse::decode(payload.as_slice()).map_err(Into::into); + } + Some(Ok(Message::Ping(payload))) => socket.send(Message::Pong(payload)).await?, + Some(Ok(Message::Pong(_))) | Some(Ok(Message::Frame(_))) => {} + Some(Ok(Message::Text(text))) => { + anyhow::bail!("public invocation returned unexpected text frame: {text}") + } + Some(Ok(Message::Close(close))) => { + anyhow::bail!("public invocation closed before its next response: {close:?}") + } + Some(Err(error)) => return Err(error.into()), + None => anyhow::bail!("public invocation connection ended before its next response"), + } + } +} + +fn public_start( + application_name: &str, + environment_name: &str, + agent_name: &str, + method_name: &str, + method_parameters: ProtoSchemaValue, +) -> PublicInvocationRequest { + let constructor_parameters = SchemaValue::Record { + fields: vec![SchemaValue::String(agent_name.to_string())], + } + .try_into() + .unwrap(); + + PublicInvocationRequest { + request: Some(public_invocation_request::Request::Start( + PublicInvocationStart { + application_name: application_name.to_string(), + environment_name: environment_name.to_string(), + agent_type_name: "StreamingRpcTarget".to_string(), + constructor_parameters: Some(constructor_parameters), + phantom_id: None, + config: Vec::new(), + method_name: method_name.to_string(), + method_parameters: Some(method_parameters), + idempotency_key: Some(IdempotencyKey::fresh().into()), + }, + )), + } +} + +fn proto_record(fields: Vec) -> ProtoSchemaValue { + SchemaValue::Record { fields }.try_into().unwrap() +} + +async fn run_public_session( + deps: &EnvBasedTestDependencies, + token: &TokenSecret, + start: PublicInvocationRequest, +) -> anyhow::Result> { + let mut socket = connect_public_invocation_socket(deps, Some(token)).await?; + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&start) + .map_err(anyhow::Error::msg)?; + send_public_request(&mut socket, &start).await?; + + let mut responses = Vec::new(); + while !state.is_complete() { + let response = receive_public_response(&mut socket).await?; + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + responses.push(response); + } + match socket.next().await { + Some(Ok(Message::Close(_))) => {} + other => anyhow::bail!("completed public invocation must close cleanly, got {other:?}"), + } + Ok(responses) +} + +fn cross_executor_agent_name( + component: &ComponentDto, + routing_table: &RoutingTable, +) -> anyhow::Result { + for index in 0..10_000 { + let name = format!("generated-streaming-rpc-cross-{index}"); + let caller = agent_id!("StreamingRpcCaller", name.clone()); + let target = agent_id!("StreamingRpcTarget", name.clone()); + let caller = AgentId::from_agent_id(component.id, &caller) + .map_err(|error| anyhow::anyhow!("invalid caller agent id: {error}"))?; + let target = AgentId::from_agent_id(component.id, &target) + .map_err(|error| anyhow::anyhow!("invalid target agent id: {error}"))?; + let caller_pod = routing_table + .lookup(&caller) + .ok_or_else(|| anyhow::anyhow!("caller agent has no executor assignment"))?; + let target_pod = routing_table + .lookup(&target) + .ok_or_else(|| anyhow::anyhow!("target agent has no executor assignment"))?; + if caller_pod != target_pod { + return Ok(name); + } + } + + anyhow::bail!("could not find caller and target agent IDs assigned to different executors") +} + +async fn invoke_agent_session( + deps: &EnvBasedTestDependencies, + component: &ComponentDto, + agent_id: &ParsedAgentId, + method_name: &str, + params: TypedSchemaValue, +) -> anyhow::Result> { + let agent_id = AgentId::from_agent_id(component.id, agent_id) + .map_err(|error| anyhow::anyhow!("invalid agent id: {error}"))?; + let (_, input) = params.into_parts(); + let input = input.try_into().map_err(anyhow::Error::msg)?; + let (frames, receiver) = mpsc::channel(8); + let request = InvocationRequest { + request: Some(invocation_request::Request::Start(InvocationStart { + agent_id: Some(agent_id.into()), + method_name: Some(method_name.to_string()), + input: Some(input), + idempotency_key: Some(IdempotencyKey::fresh().into()), + context: None, + auth_ctx: Some(AuthCtx::System.into()), + principal: None, + environment_id: None, + config: Vec::new(), + component_owner_account_id: None, + mode: golem_api_grpc::proto::golem::worker::AgentInvocationMode::Await as i32, + schedule_at: None, + freshness_disposition: + golem_api_grpc::proto::golem::worker::InvocationFreshnessDisposition::MayExist + as i32, + })), + }; + let mut state = InvocationSessionState::default(); + state + .validate_trusted_request(&request) + .map_err(anyhow::Error::msg)?; + frames.send(request).await?; + + let worker_service = deps.worker_service(); + let mut client = WorkerServiceClient::connect(format!( + "http://{}:{}", + worker_service.grpc_host(), + worker_service.gprc_port() + )) + .await?; + let mut inbound = client + .invoke_agent_session(ReceiverStream::new(receiver)) + .await? + .into_inner(); + let mut result = None; + let mut terminal = None; + while let Some(response) = inbound.message().await? { + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + match response.response { + Some(invocation_response::Response::Accepted(_)) => {} + Some(invocation_response::Response::Rejected(rejected)) => { + terminal = Some(Err(rejected.error)); + } + Some(invocation_response::Response::Result(value)) => { + if result.is_some() { + anyhow::bail!("invocation session returned more than one result"); + } + result = match value.result { + Some(invocation_session_result::Result::MethodResult(value)) => { + Some(value.try_into().map_err(anyhow::Error::msg)?) + } + Some(invocation_session_result::Result::NoResult(_)) | None => { + anyhow::bail!("invocation session returned no method result") + } + }; + } + Some(invocation_response::Response::Finished(finished)) => { + terminal = Some(match finished.outcome { + Some(invocation_session_completion::Outcome::Success(_)) => { + result.take().map(Ok).ok_or_else(|| { + anyhow::anyhow!("invocation session ended without a result") + })? + } + Some(invocation_session_completion::Outcome::Failure(failure)) => { + Err(failure.message) + } + None => anyhow::bail!("invocation session completion has no outcome"), + }); + } + Some(other) => { + anyhow::bail!("unexpected outer invocation session frame: {other:?}") + } + None => anyhow::bail!("empty outer invocation session frame"), + } + } + terminal.ok_or_else(|| anyhow::anyhow!("invocation session response ended before completion")) +} + +fn assert_streaming_report(value: SchemaValue) { + assert_eq!( + value, + SchemaValue::Record { + fields: vec![ + SchemaValue::List { + elements: vec![ + SchemaValue::U32(1), + SchemaValue::U32(2), + SchemaValue::U32(3), + ], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::U32(4), + SchemaValue::U32(5), + SchemaValue::U32(6), + ], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::U32(70), + SchemaValue::U32(80), + SchemaValue::U32(90), + ], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::String("left".to_string()), + SchemaValue::String("right".to_string()), + ], + }, + SchemaValue::List { + elements: vec![SchemaValue::U32(10), SchemaValue::U32(11)], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::String("first".to_string()), + SchemaValue::String("second".to_string()), + ], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::List { + elements: vec![SchemaValue::U32(1), SchemaValue::U32(2)], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::U32(3), + SchemaValue::U32(4), + SchemaValue::U32(5), + ], + }, + ], + }, + SchemaValue::List { + elements: vec![ + SchemaValue::String("a".to_string()), + SchemaValue::String("b".to_string()), + ], + }, + SchemaValue::List { + elements: (0..64).map(SchemaValue::U32).collect(), + }, + SchemaValue::U64(42), + ], + } + ); +} + +#[test] +#[timeout("4 minutes")] +#[tracing::instrument] +async fn generated_rust_client_streaming_rpc_cross_executor( + deps: &EnvBasedTestDependencies, +) -> anyhow::Result<()> { + let user = deps.user().await?; + let (_, environment) = user.app_and_env().await?; + let component = user + .component(&environment.id, "golem_it_agent_rpc_rust_release") + .name("golem-it:agent-rpc-rust") + .unique() + .store() + .await?; + let routing_table = deps.shard_manager().get_routing_table().await?; + let name = cross_executor_agent_name(&component, &routing_table)?; + let caller_agent_id = agent_id!("StreamingRpcCaller", name); + + let result = invoke_agent_session(deps, &component, &caller_agent_id, "run", data_value!()) + .await? + .map_err(anyhow::Error::msg)?; + assert_streaming_report(result); + + let producer_error = invoke_agent_session( + deps, + &component, + &caller_agent_id, + "call_producer_error", + data_value!(), + ) + .await? + .expect_err("producer stream error must fail the invocation session"); + assert!( + producer_error.contains("Component trapped") + || producer_error.contains("value-node index out of range: 0"), + "unexpected producer error: {producer_error}" + ); + + let stream_free_caller_id = agent_id!("StreamingRpcCaller", "stream-free-after-stream-error"); + let first = user + .invoke_and_await_agent( + &component, + &stream_free_caller_id, + "call_stream_free", + data_value!(), + ) + .await? + .into_typed::()?; + let second = user + .invoke_and_await_agent( + &component, + &stream_free_caller_id, + "call_stream_free", + data_value!(), + ) + .await? + .into_typed::()?; + assert_eq!((first, second), (1, 2)); + Ok(()) +} + +#[test] +#[timeout("4 minutes")] +#[tracing::instrument] +async fn public_websocket_invocation_forwards_scalar_and_streaming_sessions( + deps: &EnvBasedTestDependencies, +) -> anyhow::Result<()> { + let user = deps.user().await?; + let (application, environment) = user.app_and_env().await?; + user.component(&environment.id, "golem_it_agent_rpc_rust_release") + .name("golem-it:agent-rpc-rust") + .unique() + .store() + .await?; + let application_name = &application.name.0; + let environment_name = &environment.name.0; + let agent_name = format!("public-streaming-{}", uuid::Uuid::new_v4()); + + let scalar = run_public_session( + deps, + &user.token, + public_start( + application_name, + environment_name, + &agent_name, + "ping", + proto_record(Vec::new()), + ), + ) + .await?; + assert_eq!(scalar.len(), 3); + assert!(matches!( + scalar[0].response, + Some(invocation_response::Response::Accepted(_)) + )); + let Some(invocation_response::Response::Result(result)) = &scalar[1].response else { + anyhow::bail!("scalar public invocation did not return a result") + }; + let Some(invocation_session_result::Result::MethodResult(value)) = &result.result else { + anyhow::bail!("scalar public invocation returned no method value") + }; + assert_eq!( + SchemaValue::try_from(value.clone()).map_err(anyhow::Error::msg)?, + SchemaValue::U64(42) + ); + assert!(matches!( + scalar[2].response, + Some(invocation_response::Response::Finished(_)) + )); + + let produced = run_public_session( + deps, + &user.token, + public_start( + application_name, + environment_name, + &agent_name, + "produce", + proto_record(vec![SchemaValue::List { + elements: vec![ + SchemaValue::U32(3), + SchemaValue::U32(5), + SchemaValue::U32(8), + ], + }]), + ), + ) + .await?; + assert_eq!(produced.len(), 7); + let Some(invocation_response::Response::Result(result)) = &produced[1].response else { + anyhow::bail!("streaming public invocation did not return an initial result") + }; + let Some(invocation_session_result::Result::MethodResult(result_value)) = &result.result else { + anyhow::bail!("streaming public invocation returned no method value") + }; + let Some(schema_value::Value::StreamReference(stream)) = &result_value.value else { + anyhow::bail!("streaming public invocation did not return a stream reference") + }; + let items = produced[2..5] + .iter() + .map(|response| match &response.response { + Some(invocation_response::Response::OutputItem(item)) => { + assert_eq!(item.stream_id, stream.stream_id); + SchemaValue::try_from(item.value.clone().unwrap()).unwrap() + } + other => panic!("expected output item, got {other:?}"), + }) + .collect::>(); + assert_eq!( + items, + vec![ + SchemaValue::U32(3), + SchemaValue::U32(5), + SchemaValue::U32(8) + ] + ); + assert!(matches!( + produced[5].response, + Some(invocation_response::Response::OutputEnd(_)) + )); + assert!(matches!( + produced[6].response, + Some(invocation_response::Response::Finished(_)) + )); + + let input_stream_id = 71; + let stream_input = ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { + stream_id: input_stream_id, + }, + )), + }], + })), + }; + let start = public_start( + application_name, + environment_name, + &agent_name, + "consume", + stream_input, + ); + let mut socket = connect_public_invocation_socket(deps, Some(&user.token)).await?; + let mut state = InvocationSessionState::default(); + state + .validate_public_request(&start) + .map_err(anyhow::Error::msg)?; + send_public_request(&mut socket, &start).await?; + let accepted = receive_public_response(&mut socket).await?; + state + .validate_response(&accepted) + .map_err(anyhow::Error::msg)?; + + for (sequence, value) in [13_u32, 21].into_iter().enumerate() { + let request = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id: input_stream_id, + sequence: sequence as u64, + payload: Some(input_stream_item::Payload::Value( + SchemaValue::U32(value) + .try_into() + .map_err(anyhow::Error::msg)?, + )), + }, + )), + }; + state + .validate_public_request(&request) + .map_err(anyhow::Error::msg)?; + send_public_request(&mut socket, &request).await?; + let ack = receive_public_response(&mut socket).await?; + state.validate_response(&ack).map_err(anyhow::Error::msg)?; + assert!(matches!( + ack.response, + Some(invocation_response::Response::InputAck(_)) + )); + } + let end = PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputEnd( + InputStreamEnd { + stream_id: input_stream_id, + offset: 2, + }, + )), + }; + state + .validate_public_request(&end) + .map_err(anyhow::Error::msg)?; + send_public_request(&mut socket, &end).await?; + let mut consumed = None; + while !state.is_complete() { + let response = receive_public_response(&mut socket).await?; + state + .validate_response(&response) + .map_err(anyhow::Error::msg)?; + if let Some(invocation_response::Response::Result(result)) = response.response + && let Some(invocation_session_result::Result::MethodResult(value)) = result.result + { + consumed = Some(SchemaValue::try_from(value).map_err(anyhow::Error::msg)?); + } + } + assert_eq!( + consumed, + Some(SchemaValue::List { + elements: vec![SchemaValue::U32(13), SchemaValue::U32(21)], + }) + ); + + let blocked_stream_id = 73; + let blocked = public_start( + application_name, + environment_name, + &agent_name, + "consume", + ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { + stream_id: blocked_stream_id, + }, + )), + }], + })), + }, + ); + let mut blocked_socket = connect_public_invocation_socket(deps, Some(&user.token)).await?; + send_public_request(&mut blocked_socket, &blocked).await?; + let accepted = receive_public_response(&mut blocked_socket).await?; + assert!( + matches!( + accepted.response, + Some(invocation_response::Response::Accepted(_)) + ), + "blocked public invocation was not accepted: {accepted:?}" + ); + blocked_socket.send(Message::Close(None)).await?; + drop(blocked_socket); + + let subsequent = tokio::time::timeout( + std::time::Duration::from_secs(30), + run_public_session( + deps, + &user.token, + public_start( + application_name, + environment_name, + &agent_name, + "ping", + proto_record(Vec::new()), + ), + ), + ) + .await + .map_err(|_| anyhow::anyhow!("public disconnect did not cancel the blocked invocation"))??; + assert!(subsequent.iter().any(|response| matches!( + response.response, + Some(invocation_response::Response::Finished(_)) + ))); + + let capability_stream_id = 75; + let mut capability_socket = connect_public_invocation_socket(deps, Some(&user.token)).await?; + send_public_request( + &mut capability_socket, + &public_start( + application_name, + environment_name, + &agent_name, + "consume", + ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { + value: Some(schema_value::Value::StreamReference( + SchemaValueStreamReference { + stream_id: capability_stream_id, + }, + )), + }], + })), + }, + ), + ) + .await?; + let accepted = receive_public_response(&mut capability_socket).await?; + assert!(matches!( + accepted.response, + Some(invocation_response::Response::Accepted(_)) + )); + send_public_request( + &mut capability_socket, + &PublicInvocationRequest { + request: Some(public_invocation_request::Request::InputItem( + InputStreamItem { + stream_id: capability_stream_id, + sequence: 0, + payload: Some(input_stream_item::Payload::Value(ProtoSchemaValue { + value: Some(schema_value::Value::RecordValue(RecordValue { + fields: vec![ProtoSchemaValue { + value: Some(schema_value::Value::SecretValue( + SecretValue::default(), + )), + }], + })), + })), + }, + )), + }, + ) + .await?; + let close = tokio::time::timeout(std::time::Duration::from_secs(10), async { + loop { + match capability_socket.next().await { + Some(Ok(Message::Close(close))) => break Ok(close), + Some(Ok(_)) => {} + Some(Err(error)) => break Err(anyhow::Error::from(error)), + None => anyhow::bail!("capability injection ended without a protocol close"), + } + } + }) + .await + .map_err(|_| anyhow::anyhow!("capability injection did not close the public session"))??; + assert!(matches!( + close.map(|frame| frame.code), + Some(tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Protocol) + )); + drop(capability_socket); + + tokio::time::timeout( + std::time::Duration::from_secs(30), + run_public_session( + deps, + &user.token, + public_start( + application_name, + environment_name, + &agent_name, + "ping", + proto_record(Vec::new()), + ), + ), + ) + .await + .map_err(|_| anyhow::anyhow!("capability rejection did not cancel the live invocation"))??; + Ok(()) +} + +#[test] +#[timeout("2 minutes")] +#[tracing::instrument] +async fn public_websocket_invocation_enforces_auth_frames_and_rejections( + deps: &EnvBasedTestDependencies, +) -> anyhow::Result<()> { + let missing_auth = connect_public_invocation_socket(deps, None) + .await + .expect_err("missing authentication must reject the WebSocket upgrade"); + assert!(matches!( + missing_auth, + WebSocketError::Http(response) if response.status().as_u16() == 401 + )); + + let invalid_token = TokenSecret::trusted("not-a-valid-token".to_string()); + let invalid_auth = connect_public_invocation_socket(deps, Some(&invalid_token)) + .await + .expect_err("invalid authentication must reject the WebSocket upgrade"); + assert!(matches!( + invalid_auth, + WebSocketError::Http(response) if response.status().as_u16() == 401 + )); + + let user = deps.user().await?; + for invalid_message in [ + Message::Text("not protobuf".into()), + Message::Binary(vec![0xff, 0xff].into()), + ] { + let mut socket = connect_public_invocation_socket(deps, Some(&user.token)).await?; + socket.send(invalid_message).await?; + let Some(Ok(Message::Close(Some(close)))) = socket.next().await else { + anyhow::bail!("invalid public frame did not receive a close response") + }; + assert!(matches!( + close.code, + tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Unsupported + | tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Protocol + )); + } + + let rejected = run_public_session( + deps, + &user.token, + public_start( + "application-that-does-not-exist", + "environment-that-does-not-exist", + "missing-agent", + "ping", + proto_record(Vec::new()), + ), + ) + .await?; + assert_eq!(rejected.len(), 1); + let Some(invocation_response::Response::Rejected(rejected)) = &rejected[0].response else { + anyhow::bail!("unresolved public selector did not produce invocation-rejected") + }; + assert_eq!( + rejected.reason(), + InvocationRejectionReason::NotFound, + "unexpected public rejection: {}", + rejected.error + ); + Ok(()) +} diff --git a/local-run/nginx.conf b/local-run/nginx.conf index b293516379..b27c2d9466 100644 --- a/local-run/nginx.conf +++ b/local-run/nginx.conf @@ -45,6 +45,16 @@ http { proxy_pass http://worker-service; } + location = /v1/agents/invoke-agent-session { + proxy_pass http://worker-service; + proxy_http_version 1.1; + proxy_set_header Upgrade "websocket"; + proxy_set_header Connection "upgrade"; + + proxy_read_timeout 3600s; + proxy_send_timeout 3600s; + } + location /v1/agents { proxy_pass http://worker-service; } diff --git a/openapi/golem-service.yaml b/openapi/golem-service.yaml index ef34a23d51..df7de0e1dd 100644 --- a/openapi/golem-service.yaml +++ b/openapi/golem-service.yaml @@ -1989,6 +1989,72 @@ paths: security: - Cookie: [] - Token: [] + /v1/agents/invoke-agent-session: + get: + tags: + - Agent + summary: Invoke an agent through an attached live streaming session + operationId: invoke_agent_session + responses: + '101': + description: A websocket response + '400': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorsBody' + '401': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '403': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '404': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '409': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '413': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '415': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '422': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '500': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBodyWithOptionalWorkerError' + security: + - Cookie: [] + - Token: [] /v1/agents/create-agent: post: tags: diff --git a/openapi/golem-worker-service.yaml b/openapi/golem-worker-service.yaml index a99b1ad791..63b84677a2 100644 --- a/openapi/golem-worker-service.yaml +++ b/openapi/golem-worker-service.yaml @@ -1982,6 +1982,72 @@ paths: - Cookie: [] - Token: [] operationId: invoke_agent + /v1/agents/invoke-agent-session: + get: + tags: + - Agent + summary: Invoke an agent through an attached live streaming session + responses: + '101': + description: A websocket response + '400': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorsBody' + '401': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '403': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '404': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '409': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '413': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '415': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '422': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBody' + '500': + description: '' + content: + application/json; charset=utf-8: + schema: + $ref: '#/components/schemas/ErrorBodyWithOptionalWorkerError' + security: + - Cookie: [] + - Token: [] + operationId: invoke_agent_session /v1/agents/create-agent: post: tags: diff --git a/plugins/otlp-exporter.wasm b/plugins/otlp-exporter.wasm index bda3fdf87d..6fba45c92c 100644 Binary files a/plugins/otlp-exporter.wasm and b/plugins/otlp-exporter.wasm differ diff --git a/plugins/otlp-exporter/.agents/skills/golem-build/SKILL.md b/plugins/otlp-exporter/.agents/skills/golem-build/SKILL.md index ff0712be71..5635b5aebc 100644 --- a/plugins/otlp-exporter/.agents/skills/golem-build/SKILL.md +++ b/plugins/otlp-exporter/.agents/skills/golem-build/SKILL.md @@ -22,7 +22,7 @@ The build is a multi-step pipeline: 1. **Check** — verifies that required build tools are installed (e.g., `cargo` for Rust, `npm`/`node` for TypeScript). 2. **Build** — executes the build commands defined in `golem.yaml` for each component. These commands are language-specific: - **Rust**: runs `cargo build --target wasm32-wasip2` (or with `--release` for the release preset). - - **TypeScript**: runs a multi-stage pipeline — `tsc` for type checking, `golem-typegen` for metadata extraction, `rollup` for bundling, then injects the bundle into a prebuilt QuickJS WASM and optionally preinitializes it. + - **TypeScript**: runs `rollup` to bundle the component (type checking happens in-process during bundling), then injects the bundle into a prebuilt QuickJS WASM and optionally preinitializes it. - **Scala**: runs Scala.js compilation, JavaScript linking, QuickJS WASM injection, agent wrapper generation, and WASM composition. 3. **Add Metadata** — embeds component name and version into the output WASM binary. 4. **Generate Bridge** — generates bridge SDK code if the project uses inter-component communication. diff --git a/plugins/otlp-exporter/.agents/skills/golem-cloud-account-setup/SKILL.md b/plugins/otlp-exporter/.agents/skills/golem-cloud-account-setup/SKILL.md index 39c822df41..89903ca80c 100644 --- a/plugins/otlp-exporter/.agents/skills/golem-cloud-account-setup/SKILL.md +++ b/plugins/otlp-exporter/.agents/skills/golem-cloud-account-setup/SKILL.md @@ -29,7 +29,7 @@ If you need a custom cloud profile (e.g., for a different cloud endpoint): golem profile new my-cloud --url https://release.api.golem.cloud --set-active ``` -When no `--static-token` is provided, the profile uses OAuth2 (GitHub) authentication — a browser window will open on first use. +By default (equivalently `--auth oauth2`) the profile uses OAuth2 (GitHub) authentication — a browser window will open on first use. Pass `--auth static` (or `--static-token`) for token-based auth instead. ## Step 2: Authenticate @@ -72,7 +72,7 @@ golem -C account new "Team Account" "team@example.com" # Create add ## Step 4: Create and Manage API Tokens -For programmatic access (CI/CD, scripts), create static API tokens: +Create static API tokens: ```shell golem -C api-token list # List existing tokens @@ -83,10 +83,16 @@ golem -C api-token delete # Delete a t `golem api-token new` prints the token secret once, including when using structured output such as `--format json`. Store that value securely; it cannot be retrieved later. -Use a static token in a profile for non-interactive environments: +Use a static token in a profile: ```shell -golem profile new ci-cloud --url https://release.api.golem.cloud --static-token "" --set-active +golem profile new token-cloud --url https://release.api.golem.cloud --static-token "" --set-active +``` + +`--auth static` without `--static-token` prompts for the token (masked) instead of putting it on the command line: + +```shell +golem profile new my-cloud --url https://release.api.golem.cloud --auth static --set-active ``` ## Step 5: Configure Your Application for Cloud Deployment diff --git a/plugins/otlp-exporter/.agents/skills/golem-configure-mcp-server/SKILL.md b/plugins/otlp-exporter/.agents/skills/golem-configure-mcp-server/SKILL.md index 62d3c6bb2a..a545a221af 100644 --- a/plugins/otlp-exporter/.agents/skills/golem-configure-mcp-server/SKILL.md +++ b/plugins/otlp-exporter/.agents/skills/golem-configure-mcp-server/SKILL.md @@ -188,7 +188,7 @@ Constructor parameters (which identify the agent instance) are automatically inc ### Agent and Method Metadata -Add `description` and `prompt` annotations to improve MCP discoverability: +Add `description` and prompt metadata to improve MCP discoverability: **Rust:** ```rust @@ -199,9 +199,12 @@ fn increment_by(&mut self, n: u32) -> u32; **TypeScript:** ```typescript -@description("Increments the counter by n") -@prompt("Increment by a given number") -async incrementBy(n: number): Promise { ... } +incrementBy: method({ + input: { n: z.number() }, + returns: z.number(), + description: "Increments the counter by n", + promptHint: "Increment by a given number", +}), ``` **Scala:** @@ -211,7 +214,7 @@ async incrementBy(n: number): Promise { ... } def incrementBy(n: Int): Future[Int] ``` -Both annotations are optional and are included in the MCP metadata sent to clients. +Both are optional and are included in the MCP metadata sent to clients. ## Special Data Types for MCP diff --git a/plugins/otlp-exporter/.agents/skills/golem-integration-test-setup/SKILL.md b/plugins/otlp-exporter/.agents/skills/golem-integration-test-setup/SKILL.md index 45a878f9c1..a7a873cd54 100644 --- a/plugins/otlp-exporter/.agents/skills/golem-integration-test-setup/SKILL.md +++ b/plugins/otlp-exporter/.agents/skills/golem-integration-test-setup/SKILL.md @@ -236,10 +236,10 @@ secretDefaults: After the suite finishes: 1. Send `SIGINT` / `SIGTERM` to the `golem server run` process. -2. Optionally run `golem server clean --data-dir ./tests/fixtures/data` to wipe state, or just delete the directory. +2. Optionally define `localServer.dataDir: ./tests/fixtures/data` in the test manifest and run `golem server clean` to wipe state, or just delete the directory. The CLI shows the resolved path and asks for confirmation; in non-interactive teardown, use `--yes` only after confirming that the manifest points to the isolated test directory. 3. Remove `tests/fixtures/ports.json`. -Do **not** rely on `golem server clean` with the *default* data directory from a test — that would delete the developer's local development state. Always pass `--data-dir` explicitly. +Do **not** run `golem server clean -X` from a test — that would delete the developer's default local development state. Keep test data isolated with manifest `localServer.dataDir`. ## End-to-End Skeleton diff --git a/plugins/otlp-exporter/.agents/skills/golem-invoke-agent-rust/SKILL.md b/plugins/otlp-exporter/.agents/skills/golem-invoke-agent-rust/SKILL.md index d811a2a1e2..52e4f2ad47 100644 --- a/plugins/otlp-exporter/.agents/skills/golem-invoke-agent-rust/SKILL.md +++ b/plugins/otlp-exporter/.agents/skills/golem-invoke-agent-rust/SKILL.md @@ -19,7 +19,25 @@ This invokes a method on a deployed agent and **waits for the result**. The agen Text output renders return values using Rust syntax. Multiple return values are rendered as a Rust tuple, for example `(1, "ok")`. Methods returning `()` or no value print `void` in text mode. -For machine-readable output, use `--format json` or `--format yaml`. A single return value includes `result` plus `resultJson`; multiple return values include `result` plus `resultsJson`; methods returning `()` or no value omit result fields. +For machine-readable output, use `--format json`, `--format yaml`, or `--format toon`. Streaming methods emit invocation lifecycle documents in order: `accepted`, `result`, any stream `item`/terminal events, and `finished`. Scalar fields accompanying streams are in the `result` event's `value`. + +## Streaming Parameters and Results + +Use `-` as the argument for exactly one direct stream parameter to bind it to stdin. By default, `--stdin-format value` reads one Rust value per line, strips only the line terminator, preserves blank lines as values, and does not accept multiline values. `--stdin-format raw` is available only for a direct `stream` or `stream` parameter. For `stream`, each logical item is a fixed 64 KiB chunk except the final shorter chunk. For `stream`, each byte is one logical item. + +```shell +printf 'Some(1)\nNone\n' | golem agent invoke 'MyAgent()' consume_values - +cat input.bin | golem agent invoke 'MyAgent()' consume_bytes - --stdin-format raw +``` + +`--stdout-format value` is the default and renders stream items as Rust values. `--stdout-format raw` writes only bytes and requires exactly one direct `stream` or `stream` result. + +```shell +golem agent invoke 'MyAgent()' produce_values --stdout-format value +golem agent invoke 'MyAgent()' produce_bytes --stdout-format raw > output.bin +``` + +In structured CLI formats, invocation output is a sequence of lifecycle documents rather than one object or array. Streaming invocation is provisionally live-only: disconnecting or pressing Ctrl-C cancels it, and the CLI does not retry or resume the session. ## Agent ID Format @@ -73,6 +91,8 @@ golem agent invoke 'staging/MyAgent("user-123")' get_status | `-t, --trigger` | Only trigger the invocation without waiting for the result (fire-and-forget) | | `-i, --idempotency-key ` | Set a specific idempotency key; use `"-"` for auto-generated | | `--no-stream` | Disable live streaming of agent stdout/stderr/log | +| `--stdin-format value\|raw` | Select stdin stream framing; defaults to `value` | +| `--stdout-format value\|raw` | Select stream result rendering; defaults to `value` | | `--schedule-at ` | Schedule the invocation at a specific time (requires `--trigger`; ISO 8601 format) | ## Idempotency diff --git a/plugins/otlp-exporter/.agents/skills/golem-local-dev-server/SKILL.md b/plugins/otlp-exporter/.agents/skills/golem-local-dev-server/SKILL.md index 699f356e20..f1767666b8 100644 --- a/plugins/otlp-exporter/.agents/skills/golem-local-dev-server/SKILL.md +++ b/plugins/otlp-exporter/.agents/skills/golem-local-dev-server/SKILL.md @@ -86,7 +86,7 @@ When `--ports-file` is specified, the server writes a JSON file with the actual } ``` -The application manifest can mirror stable local ports with `localServer.customRequestPort` and `localServer.mcpPort`. Those fields control how deployment `subdomain` values expand: HTTP API domains resolve to `