diff --git a/docs/features/steering-and-queueing.md b/docs/features/steering-and-queueing.md index 7bfffc433d..4825068eca 100644 --- a/docs/features/steering-and-queueing.md +++ b/docs/features/steering-and-queueing.md @@ -33,6 +33,28 @@ sequenceDiagram LLM->>S: Turn completes ``` +## Message provenance + +Set the optional `source` field when forwarding a message from another agent. For example, `agent-sender-id` identifies an agent-originated message. Leave it unset for existing source-less sends. Both send and send-and-wait APIs forward the field, including with `"enqueue"` and `"immediate"` delivery. + +| SDK | Source option | +|-----|---------------| +| Node.js / TypeScript | `MessageOptions.source` | +| Python | `source=` keyword on `send` and `send_and_wait` | +| Go | `MessageOptions.Source`, a `*string` | +| .NET | `MessageOptions.Source` | +| Java | `MessageOptions.setSource(...)` | +| Rust | `MessageOptions::with_source(...)` | + +The runtime accepts `user`, `system`, `command-`, `schedule-`, and `agent-`. The SDK forwards the supplied value unchanged and omits it when unset. Rust callers using the typed RPC API can use `rpc::SendRequest::with_source(...)` without accessing generated internal fields. + +Source identifies who originated the message, while `"immediate"` mode requests delivery urgency. Neither guarantees a visible reply. Source does not change billing defaults. The runtime uses provenance for scheduling and to distinguish agent input from human authorization evidence, so derive it from trusted sender metadata rather than message text. + +A successful high-level `send` acknowledgement returns a message ID and confirms acceptance, not that the recipient has consumed the message. Do not automatically resend an accepted message merely because no reply appears. + +> [!WARNING] +> Remote backends do not necessarily preserve source end to end. Mission Control accepts it and includes it in the local echo, but drops it from remote HTTP delivery. A local source event does not prove that the remote worker received the same provenance. + ## Steering (immediate mode) Steering sends a message that is injected directly into the agent's current turn. The agent sees the message in real time and adjusts its response accordingly—useful for course-correcting without aborting the turn. diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index 8637da948c..53f1caa560 100644 --- a/dotnet/src/Session.cs +++ b/dotnet/src/Session.cs @@ -326,6 +326,7 @@ public async Task SendAsync(MessageOptions options, CancellationToken ca DisplayPrompt = options.DisplayPrompt, Attachments = options.Attachments, Mode = options.Mode, + Source = options.Source, AgentMode = options.AgentMode, Traceparent = traceparent, Tracestate = tracestate, @@ -2268,6 +2269,7 @@ internal record SendMessageRequest public string? DisplayPrompt { get; init; } public IList? Attachments { get; init; } public string? Mode { get; init; } + public string? Source { get; init; } [JsonPropertyName("agentMode")] public AgentMode? AgentMode { get; init; } public string? Traceparent { get; init; } diff --git a/dotnet/src/Types.cs b/dotnet/src/Types.cs index 9129b118e8..d3e259d359 100644 --- a/dotnet/src/Types.cs +++ b/dotnet/src/Types.cs @@ -4028,6 +4028,7 @@ private MessageOptions(MessageOptions? other) Attachments = other.Attachments is not null ? [.. other.Attachments] : null; Mode = other.Mode; + Source = other.Source; AgentMode = other.AgentMode; Prompt = other.Prompt; DisplayPrompt = other.DisplayPrompt; @@ -4050,6 +4051,16 @@ private MessageOptions(MessageOptions? other) /// public string? Mode { get; set; } /// + /// Optional provenance tag copied to the resulting user.message event. + /// Supported forms are user, system, command-<id>, + /// schedule-<numeric-id>, and agent-<id>. + /// The runtime validates the value; the SDK forwards it unchanged and omits it when null. + /// This identifies the message's origin, not a response requirement; the agent may + /// complete silently. It does not change delivery mode or billing. + /// Remote backends may not retain this metadata end-to-end. + /// + public string? Source { get; set; } + /// /// The UI mode the agent was in when this message was sent (for example "plan", "autopilot"). /// Defaults to the session's current mode when unset. /// diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index 04ef6e6405..341054ba29 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -1477,6 +1477,142 @@ public async Task Generated_Session_Rpc_Throws_When_Session_Disposed() await Assert.ThrowsAsync(() => session.Rpc.Model.GetCurrentAsync()); } + [Theory] + [InlineData(null, "enqueue", false)] + [InlineData(null, "immediate", true)] + [InlineData("user", "enqueue", false)] + [InlineData("system", "immediate", true)] + [InlineData("command-review", "immediate", false)] + [InlineData("schedule-123", "enqueue", true)] + [InlineData("agent-reviewer", "enqueue", false)] + [InlineData("agent-reviewer", "immediate", true)] + [InlineData("", "enqueue", false)] + [InlineData("future-source", "immediate", true)] + public async Task Send_Source_Is_Serialized_Without_Changing_Other_Fields( + string? source, string mode, bool waitForReply) + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + await using var session = await client.CreateSessionAsync(new SessionConfig()); + using var activity = new Activity("send-source").SetIdFormat(ActivityIdFormat.W3C).Start(); + activity.TraceStateString = "test=value"; + var options = new MessageOptions + { + Prompt = "real prompt", + DisplayPrompt = "display prompt", + Source = source, + Mode = mode, + AgentMode = AgentMode.Plan, + Attachments = [new AttachmentFile { Path = "/test.txt", DisplayName = "test.txt" }], + RequestHeaders = new Dictionary { ["X-Request-ID"] = "request-1" } + }.Clone(); + + if (waitForReply) + { + var replyTask = session.SendAndWaitAsync(options, TimeSpan.FromSeconds(5)); + await WaitForRequestAsync(server, "session.send"); + await server.SendEventAsync(session.SessionId, new AssistantMessageEvent + { + Id = Guid.NewGuid(), + Data = new AssistantMessageData { Content = "reply", MessageId = "assistant-1" } + }); + await server.SendEventAsync(session.SessionId, new SessionIdleEvent + { + Id = Guid.NewGuid(), + Data = new SessionIdleData { Mode = SessionMode.Interactive } + }); + var reply = await replyTask; + Assert.NotNull(reply); + Assert.Equal("reply", reply.Data.Content); + } + else + { + Assert.Equal("message-1", await session.SendAsync(options)); + } + + var request = Assert.Single(server.Requests, request => request.Method == "session.send").Params; + if (source is null) + { + Assert.False(request.TryGetProperty("source", out _)); + } + else + { + Assert.Equal(source, request.GetProperty("source").GetString()); + } + Assert.Equal(session.SessionId, request.GetProperty("sessionId").GetString()); + Assert.Equal("real prompt", request.GetProperty("prompt").GetString()); + Assert.Equal("display prompt", request.GetProperty("displayPrompt").GetString()); + Assert.Equal(mode, request.GetProperty("mode").GetString()); + Assert.Equal("plan", request.GetProperty("agentMode").GetString()); + var attachment = Assert.Single(request.GetProperty("attachments").EnumerateArray()); + Assert.Equal("file", attachment.GetProperty("type").GetString()); + Assert.Equal("/test.txt", attachment.GetProperty("path").GetString()); + Assert.Equal("test.txt", attachment.GetProperty("displayName").GetString()); + Assert.Equal("request-1", request.GetProperty("requestHeaders").GetProperty("X-Request-ID").GetString()); + Assert.Equal(activity.Id, request.GetProperty("traceparent").GetString()); + Assert.Equal("test=value", request.GetProperty("tracestate").GetString()); + Assert.False(request.TryGetProperty("billable", out _)); + Assert.False(request.TryGetProperty("wait", out _)); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Send_String_Overloads_Omit_Source_And_Keep_Defaults(bool waitForReply) + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + await using var session = await client.CreateSessionAsync(new SessionConfig()); + + if (waitForReply) + { + var replyTask = session.SendAndWaitAsync("prompt", TimeSpan.FromSeconds(5)); + await WaitForRequestAsync(server, "session.send"); + await server.SendEventAsync(session.SessionId, new SessionIdleEvent + { + Id = Guid.NewGuid(), + Data = new SessionIdleData { Mode = SessionMode.Interactive } + }); + Assert.Null(await replyTask); + } + else + { + Assert.Equal("message-1", await session.SendAsync("prompt")); + } + + var request = Assert.Single(server.Requests, request => request.Method == "session.send").Params; + Assert.Equal("prompt", request.GetProperty("prompt").GetString()); + Assert.False(request.TryGetProperty("source", out _)); + Assert.False(request.TryGetProperty("mode", out _)); + Assert.False(request.TryGetProperty("agentMode", out _)); + Assert.False(request.TryGetProperty("billable", out _)); + } + + [Theory] + [InlineData(null)] + [InlineData("agent-reviewer")] + public async Task Rpc_Send_Source_Preserves_Explicit_Billing(string? source) + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + await using var session = await client.CreateSessionAsync(new SessionConfig()); + + var result = await session.Rpc.SendAsync("prompt", source: source, mode: SendMode.Immediate, billable: false); + + Assert.Equal("message-1", result.MessageId); + var request = Assert.Single(server.Requests, request => request.Method == "session.send").Params; + Assert.Equal("immediate", request.GetProperty("mode").GetString()); + Assert.False(request.GetProperty("billable").GetBoolean()); + if (source is null) + { + Assert.False(request.TryGetProperty("source", out _)); + } + else + { + Assert.Equal(source, request.GetProperty("source").GetString()); + } + } + [Fact] public async Task SendAndWaitAsync_Skips_Autopilot_Continuation_Idle() { @@ -2031,6 +2167,22 @@ public void CloseConnection() _stream?.Dispose(); } + public async Task SendEventAsync(string sessionId, SessionEvent evt) + { + var stream = _stream ?? throw new InvalidOperationException("Client is not connected."); + using var document = JsonDocument.Parse(evt.ToJson()); + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["method"] = "session.event", + ["params"] = new Dictionary + { + ["sessionId"] = sessionId, + ["event"] = document.RootElement + } + }, _cts.Token); + } + public async Task SendRequestAsync(string method, Dictionary parameters) { var stream = _stream ?? throw new InvalidOperationException("Client is not connected."); diff --git a/dotnet/test/Unit/CloneTests.cs b/dotnet/test/Unit/CloneTests.cs index 2f213525f6..cdd8863619 100644 --- a/dotnet/test/Unit/CloneTests.cs +++ b/dotnet/test/Unit/CloneTests.cs @@ -275,20 +275,31 @@ public void ResumeSessionConfig_Clone_PreservesMcpServersComparer() Assert.True(clone.McpServers!.ContainsKey("SERVER")); } - [Fact] - public void MessageOptions_Clone_CopiesAllProperties() + [Theory] + [InlineData(null)] + [InlineData("agent-reviewer")] + public void MessageOptions_Clone_CopiesAllProperties(string? source) { var original = new MessageOptions { Prompt = "Hello", Attachments = [new AttachmentFile { Path = "/test.txt", DisplayName = "test.txt" }], - Mode = "chat", + Mode = "immediate", + Source = source, + AgentMode = AgentMode.Plan, + DisplayPrompt = "Display hello", + RequestHeaders = new Dictionary { ["X-Request-ID"] = "request-1" }, }; var clone = original.Clone(); Assert.Equal(original.Prompt, clone.Prompt); Assert.Equal(original.Mode, clone.Mode); + Assert.Equal(original.Source, clone.Source); + Assert.Equal(original.AgentMode, clone.AgentMode); + Assert.Equal(original.DisplayPrompt, clone.DisplayPrompt); + Assert.Equal(original.RequestHeaders, clone.RequestHeaders); + Assert.NotSame(original.RequestHeaders, clone.RequestHeaders); Assert.Single(clone.Attachments!); } diff --git a/go/session.go b/go/session.go index f4bb0d38c7..2dd2a8a9a1 100644 --- a/go/session.go +++ b/go/session.go @@ -432,6 +432,7 @@ func (s *Session) Send(ctx context.Context, options MessageOptions) (string, err req := sessionSendRequest{ SessionID: s.SessionID, Prompt: options.Prompt, + Source: options.Source, DisplayPrompt: options.DisplayPrompt, Attachments: options.Attachments, Mode: options.Mode, diff --git a/go/session_test.go b/go/session_test.go index bdf14887a3..d5305ff0cc 100644 --- a/go/session_test.go +++ b/go/session_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "io" + "reflect" "strconv" "strings" "sync" @@ -95,6 +96,51 @@ func ptr[T any](value T) *T { return &value } +func TestSession_SendPreservesSource(t *testing.T) { + for _, mode := range []string{"", "enqueue", "immediate"} { + for _, source := range []*string{nil, ptr("agent-sender-id")} { + t.Run(fmt.Sprintf("%s/source=%v", mode, source != nil), func(t *testing.T) { + options := MessageOptions{ + Prompt: "Agent update", + Source: source, + Mode: mode, + AgentMode: AgentModePlan, + DisplayPrompt: "Update from sender", + RequestHeaders: map[string]string{"X-Custom-Tag": "value-1"}, + Attachments: []Attachment{ + &AttachmentFile{Path: "report.txt", DisplayName: "Report"}, + }, + } + params := captureModelRequest(t, "session.send", func(session *Session) error { + ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second) + defer cancel() + _, err := session.Send(ctx, options) + return err + }) + expected := map[string]any{ + "sessionId": "session-1", + "prompt": "Agent update", + "agentMode": "plan", + "displayPrompt": "Update from sender", + "requestHeaders": map[string]any{"X-Custom-Tag": "value-1"}, + "attachments": []any{ + map[string]any{"type": "file", "path": "report.txt", "displayName": "Report"}, + }, + } + if mode != "" { + expected["mode"] = mode + } + if source != nil { + expected["source"] = *source + } + if !reflect.DeepEqual(params, expected) { + t.Fatalf("got params %#v, want %#v", params, expected) + } + }) + } + } +} + func TestSession_SetModelForwardsContextTier(t *testing.T) { tier := ContextTierLongContext params := captureSetModelRequest(t, &SetModelOptions{ContextTier: &tier}) @@ -513,6 +559,15 @@ func readTestJSONRPCFrame(r io.Reader) ([]byte, error) { } func TestSession_SendAndWaitSkipsAutopilotContinuationIdle(t *testing.T) { + for _, source := range []*string{nil, ptr("agent-sender-id")} { + t.Run(fmt.Sprintf("source=%v", source != nil), func(t *testing.T) { + testSessionSendAndWaitSkipsAutopilotContinuationIdle(t, source) + }) + } +} + +func testSessionSendAndWaitSkipsAutopilotContinuationIdle(t *testing.T, source *string) { + t.Helper() stdinR, stdinW := io.Pipe() stdoutR, stdoutW := io.Pipe() defer stdinR.Close() @@ -536,6 +591,7 @@ func TestSession_SendAndWaitSkipsAutopilotContinuationIdle(t *testing.T) { var request struct { ID json.RawMessage `json:"id"` Method string `json:"method"` + Params map[string]any `json:"params"` } if err := json.Unmarshal(frame, &request); err != nil { errCh <- err @@ -545,6 +601,14 @@ func TestSession_SendAndWaitSkipsAutopilotContinuationIdle(t *testing.T) { errCh <- fmt.Errorf("expected session.send, got %s", request.Method) return } + if source != nil && request.Params["source"] != *source { + errCh <- fmt.Errorf("expected agent source, got %v", request.Params["source"]) + return + } + if _, present := request.Params["source"]; source == nil && present { + errCh <- fmt.Errorf("expected omitted source, got %v", request.Params["source"]) + return + } response := map[string]any{ "jsonrpc": "2.0", @@ -576,7 +640,9 @@ func TestSession_SendAndWaitSkipsAutopilotContinuationIdle(t *testing.T) { resultCh := make(chan *SessionEvent, 1) go func() { - result, err := session.SendAndWait(t.Context(), MessageOptions{Prompt: "keep going"}) + result, err := session.SendAndWait(t.Context(), MessageOptions{ + Prompt: "keep going", Source: source, + }) if err != nil { errCh <- err return diff --git a/go/types.go b/go/types.go index 4a45a4019c..4c9f583cf8 100644 --- a/go/types.go +++ b/go/types.go @@ -2435,6 +2435,12 @@ type ToolBinaryResult struct { type MessageOptions struct { // Prompt is the message to send Prompt string + // Source is optional message provenance, omitted when nil. The runtime accepts + // "user", "system", "command-", "schedule-", and "agent-". + // This is not a response requirement: agent messages may complete without a + // visible reply. It does not control billing, and remote backends may not + // preserve it end to end. + Source *string // Attachments are file or directory attachments Attachments []Attachment // Mode is the message delivery mode (default: "enqueue") @@ -2913,6 +2919,7 @@ type sessionAbortRequest struct { type sessionSendRequest struct { SessionID string `json:"sessionId"` Prompt string `json:"prompt"` + Source *string `json:"source,omitempty"` DisplayPrompt string `json:"displayPrompt,omitempty"` Attachments []Attachment `json:"attachments,omitempty"` Mode string `json:"mode,omitempty"` diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java index c072e31ec1..6c93919c9b 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java @@ -564,6 +564,7 @@ public CompletableFuture send(MessageOptions options) { request.setAgentMode(options.getAgentMode()); request.setRequestHeaders(options.getRequestHeaders()); request.setDisplayPrompt(options.getDisplayPrompt()); + request.setSource(options.getSource()); return rpc.invoke("session.send", request, SendMessageResponse.class).thenApply(SendMessageResponse::messageId); } diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/MessageOptions.java b/java/sdk/src/main/java/com/github/copilot/rpc/MessageOptions.java index c781011ff8..c72bf92298 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/MessageOptions.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/MessageOptions.java @@ -48,6 +48,7 @@ public class MessageOptions { private AgentMode agentMode; private Map requestHeaders; private String displayPrompt; + private String source; /** * Gets the message prompt. @@ -201,6 +202,37 @@ public MessageOptions setDisplayPrompt(String displayPrompt) { return this; } + /** + * Gets the optional message provenance tag. + * + * @return the source, or {@code null} if not set + */ + public String getSource() { + return source; + } + + /** + * Sets the optional provenance tag for the resulting {@code user.message} + * event. + *

+ * Accepted forms are {@code user}, {@code system}, {@code command-}, + * {@code schedule-}, and {@code agent-}. The SDK forwards + * strings unchanged. The default is {@code null}, which omits the field from + * the request. + *

+ * Source describes where a message originated, not whether it requires a + * response; the agent may complete silently. It does not change billing or + * delivery mode. Remote backends may not retain it end-to-end. + * + * @param source + * the provenance tag, or {@code null} to omit it + * @return this options instance for method chaining + */ + public MessageOptions setSource(String source) { + this.source = source; + return this; + } + /** * Creates a shallow clone of this {@code MessageOptions} instance. *

@@ -220,6 +252,7 @@ public MessageOptions clone() { copy.agentMode = this.agentMode; copy.requestHeaders = this.requestHeaders != null ? new HashMap<>(this.requestHeaders) : null; copy.displayPrompt = this.displayPrompt; + copy.source = this.source; return copy; } diff --git a/java/sdk/src/main/java/com/github/copilot/rpc/SendMessageRequest.java b/java/sdk/src/main/java/com/github/copilot/rpc/SendMessageRequest.java index c87dda7623..a977d85862 100644 --- a/java/sdk/src/main/java/com/github/copilot/rpc/SendMessageRequest.java +++ b/java/sdk/src/main/java/com/github/copilot/rpc/SendMessageRequest.java @@ -46,6 +46,9 @@ public final class SendMessageRequest { @JsonProperty("displayPrompt") private String displayPrompt; + @JsonProperty("source") + private String source; + /** Gets the session ID. @return the session ID */ public String getSessionId() { return sessionId; @@ -120,4 +123,16 @@ public String getDisplayPrompt() { public void setDisplayPrompt(String displayPrompt) { this.displayPrompt = displayPrompt; } + + /** Gets the provenance tag. @return the source, or {@code null} if not set */ + public String getSource() { + return source; + } + + /** + * Sets the provenance tag. @param source the source, or {@code null} to omit it + */ + public void setSource(String source) { + this.source = source; + } } diff --git a/java/sdk/src/test/java/com/github/copilot/MessageSourceTest.java b/java/sdk/src/test/java/com/github/copilot/MessageSourceTest.java new file mode 100644 index 0000000000..d8450ecf0c --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/MessageSourceTest.java @@ -0,0 +1,166 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +import java.io.IOException; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; + +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.NullSource; +import org.junit.jupiter.params.provider.ValueSource; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.github.copilot.rpc.AgentMode; +import com.github.copilot.rpc.Attachment; +import com.github.copilot.rpc.CopilotClientOptions; +import com.github.copilot.rpc.MessageOptions; +import com.github.copilot.rpc.PermissionHandler; +import com.github.copilot.rpc.SessionConfig; + +class MessageSourceTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @ParameterizedTest + @NullSource + @ValueSource(strings = {"user", "system", "command-review", "schedule-42", "agent-helper", " custom ", ""}) + void sendPreservesSourceAndOtherOptions(String source) throws Exception { + try (var server = new MessageServer(); + var client = new CopilotClient(new CopilotClientOptions().setCliUrl(server.url())); + var session = client.createSession(new SessionConfig().setSessionId("source-test") + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL)).get(5, TimeUnit.SECONDS)) { + for (String mode : List.of("enqueue", "immediate")) { + var options = options(mode); + assertNull(options.getSource()); + if (source != null) { + assertSame(options, options.setSource(source)); + } + assertEquals("message-id", session.send(options).get(5, TimeUnit.SECONDS)); + assertRequest(server, options, source); + } + } + } + + @ParameterizedTest + @NullSource + @ValueSource(strings = {"user", "system", "command-review", "schedule-42", "agent-helper", " custom ", ""}) + void sendAndWaitPreservesClonedSourceAndOtherOptions(String source) throws Exception { + try (var server = new MessageServer(); + var client = new CopilotClient(new CopilotClientOptions().setCliUrl(server.url())); + var session = client.createSession(new SessionConfig().setSessionId("source-test") + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL)).get(5, TimeUnit.SECONDS)) { + for (String mode : List.of("enqueue", "immediate")) { + var original = options(mode).setSource(source); + var copy = original.clone(); + assertEquals(MAPPER.valueToTree(original), MAPPER.valueToTree(copy)); + original.setSource("agent-other"); + assertEquals(source, copy.getSource()); + + assertNull(session.sendAndWait(copy, 5000).get(5, TimeUnit.SECONDS)); + assertRequest(server, copy, source); + } + } + } + + private static MessageOptions options(String mode) { + return new MessageOptions().setPrompt("Review this file").setDisplayPrompt("Review") + .setAttachments(List.of(new Attachment("file", "/workspace/main.java", "Main"))).setMode(mode) + .setAgentMode(AgentMode.PLAN).setRequestHeaders(Map.of("x-request-id", "trace-123")); + } + + private static void assertRequest(MessageServer server, MessageOptions options, String source) throws Exception { + JsonNode actual = server.requests.poll(5, TimeUnit.SECONDS); + ObjectNode expected = MAPPER.valueToTree(options); + expected.put("sessionId", "source-test"); + assertEquals(expected, actual); + if (source == null) { + assertFalse(actual.has("source")); + } else { + assertEquals(source, actual.path("source").asText()); + } + assertFalse(actual.has("billable"), "Provenance must not override the default billing behavior"); + } + + /** Uses the same loopback JSON-RPC setup as BuiltinPluginDirectoriesTest. */ + private static final class MessageServer implements AutoCloseable { + + private final ServerSocket listener; + private final Thread acceptThread; + private final CompletableFuture ready = new CompletableFuture<>(); + private final LinkedBlockingQueue requests = new LinkedBlockingQueue<>(); + + MessageServer() throws IOException { + listener = new ServerSocket(0, 1, InetAddress.getByName("127.0.0.1")); + acceptThread = new Thread(this::accept, "message-source-runtime"); + acceptThread.setDaemon(true); + acceptThread.start(); + } + + String url() { + return "127.0.0.1:" + listener.getLocalPort(); + } + + private void accept() { + try { + Socket socket = listener.accept(); + ready.complete(JsonRpcClient.fromSocket(socket, this::registerHandlers)); + } catch (IOException e) { + ready.completeExceptionally(e); + } + } + + private void registerHandlers(JsonRpcClient rpc) { + rpc.registerMethodHandler("connect", + (id, params) -> respond(rpc, id, Map.of("ok", true, "protocolVersion", 3, "version", "test"))); + rpc.registerMethodHandler("session.create", + (id, params) -> respond(rpc, id, Map.of("sessionId", "source-test"))); + rpc.registerMethodHandler("session.send", (id, params) -> { + requests.add(params); + respond(rpc, id, Map.of("messageId", "message-id")); + try { + rpc.notify("session.event", + Map.of("sessionId", "source-test", "event", Map.of("id", UUID.randomUUID().toString(), + "type", "session.idle", "timestamp", "2026-01-01T00:00:00Z", "data", Map.of()))); + } catch (IOException e) { + throw new java.io.UncheckedIOException(e); + } + }); + rpc.registerMethodHandler("session.detach", (id, params) -> respond(rpc, id, Map.of("success", true))); + } + + private static void respond(JsonRpcClient rpc, String id, Object result) { + try { + rpc.sendResponse(id, result); + } catch (IOException e) { + throw new java.io.UncheckedIOException(e); + } + } + + @Override + public void close() throws Exception { + listener.close(); + acceptThread.join(5000); + JsonRpcClient rpc = ready.getNow(null); + if (rpc != null) { + rpc.close(); + } + } + } +} diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index b7fc7837a5..4c2be14299 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -719,6 +719,7 @@ export class CopilotSession { ...(await getTraceContext(this.traceContextProvider)), sessionId: this.sessionId, prompt: options.prompt, + source: options.source, displayPrompt: options.displayPrompt, attachments: options.attachments, mode: options.mode, diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts index 0f15749f71..a2add86d39 100644 --- a/nodejs/src/types.ts +++ b/nodejs/src/types.ts @@ -3323,6 +3323,14 @@ export interface MessageOptions { */ prompt: string; + /** + * Optional message provenance, omitted when unset. The runtime accepts `user`, `system`, + * `command-`, `schedule-`, and `agent-`. + * This is not a response requirement: agent messages may complete without a visible reply. + * It does not control billing, and remote backends may not preserve it end to end. + */ + source?: string; + /** * File, directory, selection, or blob attachments */ diff --git a/nodejs/test/session-source.test.ts b/nodejs/test/session-source.test.ts new file mode 100644 index 0000000000..3bea80f0cd --- /dev/null +++ b/nodejs/test/session-source.test.ts @@ -0,0 +1,67 @@ +import { PassThrough } from "node:stream"; +import { describe, expect, it, onTestFinished } from "vitest"; +import { createMessageConnection } from "vscode-jsonrpc/node.js"; +import { CopilotSession } from "../src/session.js"; +import type { MessageOptions } from "../src/types.js"; + +describe.each(["send", "sendAndWait"] as const)("%s message provenance", (method) => { + for (const mode of [undefined, "enqueue", "immediate"] as const) { + it.each([undefined, "agent-sender-id"])( + `preserves source %s and other options with mode ${mode}`, + async (source) => { + const requests = new PassThrough(); + const responses = new PassThrough(); + const connection = createMessageConnection(responses, requests); + const server = createMessageConnection(requests, responses); + const session = new CopilotSession("session-1", connection); + onTestFinished(() => { + connection.dispose(); + server.dispose(); + requests.destroy(); + responses.destroy(); + }); + + const options: MessageOptions = { + prompt: "Agent update", + source, + mode, + agentMode: "plan", + attachments: [{ type: "file", path: "report.txt", displayName: "Report" }], + displayPrompt: "Update from sender", + requestHeaders: { "X-Custom-Tag": "value-1" }, + }; + let received: Record | undefined; + server.onRequest("session.send", (params: Record) => { + received = params; + session._dispatchEvent({ + type: "session.idle", + id: "00000000-0000-4000-8000-000000000001", + parentId: null, + timestamp: new Date().toISOString(), + ephemeral: true, + data: {}, + }); + return { messageId: "message-1" }; + }); + connection.listen(); + server.listen(); + + const result = await session[method](options); + expect(result).toBe(method === "send" ? "message-1" : undefined); + expect(received).toEqual({ + sessionId: "session-1", + prompt: options.prompt, + agentMode: options.agentMode, + attachments: options.attachments, + displayPrompt: options.displayPrompt, + requestHeaders: options.requestHeaders, + ...(mode === undefined ? {} : { mode }), + ...(source === undefined ? {} : { source }), + }); + if (source === undefined) { + expect(received).not.toHaveProperty("source"); + } + } + ); + } +}); diff --git a/python/copilot/session.py b/python/copilot/session.py index b5823923ee..4b6e0e8800 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -1710,6 +1710,7 @@ async def send( self, prompt: str, *, + source: str | None = None, attachments: list[Attachment] | None = None, mode: Literal["enqueue", "immediate"] | None = None, agent_mode: Literal["interactive", "plan", "autopilot", "shell"] | None = None, @@ -1725,6 +1726,11 @@ async def send( Args: prompt: The message text to send. + source: Optional message provenance, omitted when None. The runtime accepts + ``user``, ``system``, ``command-``, ``schedule-``, and + ``agent-``. This is not a response requirement: agent messages may + complete without a visible reply. It does not control billing, and + remote backends may not preserve it end to end. attachments: Optional file, directory, or selection attachments. mode: Message delivery mode (``"enqueue"`` or ``"immediate"``). agent_mode: The UI mode the agent was in when this message was sent @@ -1750,6 +1756,8 @@ async def send( "sessionId": self.session_id, "prompt": prompt, } + if source is not None: + params["source"] = source if attachments is not None: params["attachments"] = attachments if mode is not None: @@ -1779,6 +1787,7 @@ async def send_and_wait( self, prompt: str, *, + source: str | None = None, attachments: list[Attachment] | None = None, mode: Literal["enqueue", "immediate"] | None = None, agent_mode: Literal["interactive", "plan", "autopilot", "shell"] | None = None, @@ -1797,6 +1806,9 @@ async def send_and_wait( Args: prompt: The message text to send. + source: Optional message provenance, forwarded unchanged to :meth:`send`. + It does not require a visible reply or control billing. Remote + backends may not preserve it end to end. attachments: Optional file, directory, or selection attachments. mode: Message delivery mode (``"enqueue"`` or ``"immediate"``). agent_mode: The UI mode the agent was in when this message was sent @@ -1860,6 +1872,7 @@ def handler(event: SessionEventTypeAlias) -> None: try: await self.send( prompt, + source=source, attachments=attachments, mode=mode, agent_mode=agent_mode, diff --git a/python/test_session.py b/python/test_session.py index d58ce94091..6a7294656b 100644 --- a/python/test_session.py +++ b/python/test_session.py @@ -7,7 +7,7 @@ import pytest -from copilot.session import CopilotSession +from copilot.session import Attachment, CopilotSession from copilot.session_events import ( AssistantMessageData, ExternalToolCompletedData, @@ -29,6 +29,60 @@ def _event(data, event_type: SessionEventType) -> SessionEvent: ) +@pytest.mark.parametrize("wait", [False, True]) +@pytest.mark.parametrize("mode", [None, "enqueue", "immediate"]) +@pytest.mark.parametrize("source", [None, "agent-sender-id"]) +@pytest.mark.asyncio +async def test_send_preserves_source_and_other_options(wait, mode, source): + client = Mock() + session = CopilotSession("session-1", client) + + async def respond(_method, _params): + session._dispatch_event(_event(SessionIdleData(), SessionEventType.SESSION_IDLE)) + return {"messageId": "message-1"} + + client.request = AsyncMock(side_effect=respond) + attachments: list[Attachment] = [ + {"type": "file", "path": "report.txt", "displayName": "Report"} + ] + send = session.send_and_wait if wait else session.send + result = await send( + "Agent update", + source=source, + mode=mode, + agent_mode="plan", + attachments=attachments, + display_prompt="Update from sender", + request_headers={"X-Custom-Tag": "value-1"}, + ) + + expected = { + "sessionId": "session-1", + "prompt": "Agent update", + "agentMode": "plan", + "attachments": attachments, + "displayPrompt": "Update from sender", + "requestHeaders": {"X-Custom-Tag": "value-1"}, + } + if mode is not None: + expected["mode"] = mode + if source is not None: + expected["source"] = source + client.request.assert_awaited_once_with("session.send", expected) + assert result == (None if wait else "message-1") + + +@pytest.mark.asyncio +async def test_send_default_omits_source(): + client = Mock() + client.request = AsyncMock(return_value={"messageId": "message-1"}) + session = CopilotSession("session-1", client) + assert await session.send("Human message") == "message-1" + client.request.assert_awaited_once_with( + "session.send", {"sessionId": "session-1", "prompt": "Human message"} + ) + + @pytest.mark.asyncio async def test_send_and_wait_skips_autopilot_continuation_idle(): client = Mock() diff --git a/rust/src/rpc.rs b/rust/src/rpc.rs index a08a501cb2..b9b40f8ee9 100644 --- a/rust/src/rpc.rs +++ b/rust/src/rpc.rs @@ -1,6 +1,7 @@ //! JSON-RPC request/response types and typed namespace builders. //! -//! All types are auto-generated from the Copilot CLI protocol schemas. +//! Types are auto-generated from the Copilot CLI protocol schemas, with +//! handwritten convenience methods where needed. //! This module is the stable public access point — the underlying //! crate-private modules where the types are defined are an //! implementation detail whose layout may change. @@ -10,3 +11,14 @@ pub use crate::generated::api_types::*; pub use crate::generated::rpc::*; + +impl SendRequest { + /// Set the message provenance tag without changing billing or delivery options. + /// + /// See [`crate::types::MessageOptions::source`] for accepted forms and + /// remote-delivery limitations. When unset, the tag is omitted. + pub fn with_source(mut self, source: impl Into) -> Self { + self.source = Some(source.into()); + self + } +} diff --git a/rust/src/session.rs b/rust/src/session.rs index 0e64d6061c..82e97b6606 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -522,6 +522,9 @@ impl Session { "sessionId": self.id, "prompt": opts.prompt, }); + if let Some(source) = opts.source { + params["source"] = serde_json::to_value(source)?; + } if let Some(m) = opts.mode { params["mode"] = serde_json::to_value(m)?; } diff --git a/rust/src/types.rs b/rust/src/types.rs index a99f00a19f..43c2c9cfd8 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -5380,6 +5380,16 @@ pub enum AgentMode { pub struct MessageOptions { /// The user prompt to send. pub prompt: String, + /// Optional provenance tag identifying the message origin. + /// + /// The runtime accepts `user`, `system`, `command-`, + /// `schedule-`, and `agent-`. The SDK forwards + /// the value unchanged and omits it when `None` (the default). + /// + /// Provenance does not specify whether a reply is required: agent messages + /// may complete without a visible reply. It does not control billing. + /// Remote backends may not preserve this tag across delivery. + pub source: Option, /// Optional message delivery mode for this turn. /// /// Controls whether the prompt is queued behind in-flight work @@ -5419,6 +5429,7 @@ impl MessageOptions { pub fn new(prompt: impl Into) -> Self { Self { prompt: prompt.into(), + source: None, mode: None, agent_mode: None, attachments: None, @@ -5430,6 +5441,12 @@ impl MessageOptions { } } + /// Set the message provenance tag. See [`Self::source`] for accepted forms. + pub fn with_source(mut self, source: impl Into) -> Self { + self.source = Some(source.into()); + self + } + /// Set the message delivery mode for this turn. /// /// Pass [`DeliveryMode::Immediate`] to interrupt the session and run diff --git a/rust/tests/api_types_test.rs b/rust/tests/api_types_test.rs index 9429a2bb6f..4f0c37913b 100644 --- a/rust/tests/api_types_test.rs +++ b/rust/tests/api_types_test.rs @@ -7,7 +7,7 @@ use github_copilot_sdk::rpc::{ Extension, ExtensionList, ExtensionSource, ExtensionStatus, ExtensionsDisableRequest, ExtensionsEnableRequest, FleetStartRequest, FleetStartResult, ModelSwitchAutoTierRequest, ModelSwitchAutoTierResult, ModelSwitchAutoTierStatus, QueuePendingItems, QueuePendingItemsKind, - SandboxConfig, SendAgentMode, TasksStartAgentRequest, + SandboxConfig, SendAgentMode, SendMode, SendRequest, TasksStartAgentRequest, }; use github_copilot_sdk::session_events::{ PermissionRequest, PermissionRequestedData, SessionEventData, TypedSessionEvent, @@ -115,6 +115,46 @@ fn fleet_start_request_and_result_fields_are_accessible() { assert!(result.started); } +#[test] +fn send_request_source_builder_changes_only_provenance() { + let mut request = SendRequest::default(); + request.prompt = "Agent update".to_string(); + request.attachments = Some(vec![serde_json::json!({ + "type": "file", + "path": "report.txt", + "displayName": "Report" + })]); + request.display_prompt = Some("Update from sender".to_string()); + request.mode = Some(SendMode::Immediate); + request.agent_mode = Some(SendAgentMode::Plan); + request.billable = Some(false); + request.prepend = Some(true); + request.required_tool = Some("read_file".to_string()); + request.request_headers = Some( + [("X-Custom-Tag".to_string(), "value-1".to_string())] + .into_iter() + .collect(), + ); + request.traceparent = Some("00-trace-parent-01".to_string()); + request.tracestate = Some("vendor=value".to_string()); + request.wait = Some(true); + + let mut expected = serde_json::to_value(&request).unwrap(); + assert!(expected.get("source").is_none()); + expected["source"] = serde_json::json!("agent-sender-id"); + + let sourced = request.with_source("agent-sender-id".to_string()); + assert_eq!(serde_json::to_value(sourced).unwrap(), expected); +} + +#[test] +fn send_request_default_omits_source_and_billing() { + assert_eq!( + serde_json::to_value(SendRequest::default()).unwrap(), + serde_json::json!({ "prompt": "" }) + ); +} + #[test] fn tasks_start_agent_request_fields_are_accessible() { let request = TasksStartAgentRequest { diff --git a/rust/tests/e2e/copilot_request_handler.rs b/rust/tests/e2e/copilot_request_handler.rs index 478845f48e..54884fc667 100644 --- a/rust/tests/e2e/copilot_request_handler.rs +++ b/rust/tests/e2e/copilot_request_handler.rs @@ -24,12 +24,13 @@ use async_trait::async_trait; use bytes::Bytes; use futures_util::{SinkExt, StreamExt}; use github_copilot_sdk::handler::ApproveAllHandler; -use github_copilot_sdk::session_events::AssistantMessageData; +use github_copilot_sdk::rpc::{SendMode, SendRequest}; +use github_copilot_sdk::session_events::{AssistantMessageData, UserMessageData}; use github_copilot_sdk::{ CopilotHttpRequest, CopilotHttpResponse, CopilotRequestContext, CopilotRequestError, CopilotRequestHandler, CopilotWebSocketForwarder, CopilotWebSocketHandler, - CopilotWebSocketResponse, MessageOptions, ProviderConfig, SessionConfig, SessionEvent, - forward_http, + CopilotWebSocketResponse, DeliveryMode, MessageOptions, ProviderConfig, SessionConfig, + SessionEvent, forward_http, }; use http::header::{HeaderName, HeaderValue}; use http::{HeaderMap, Uri}; @@ -38,7 +39,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio_tungstenite::tungstenite::Message; -use super::support::with_e2e_context_no_snapshot; +use super::support::{collect_until_idle, with_e2e_context_no_snapshot}; const SYNTHETIC_TEXT: &str = "OK from the synthetic stream."; const HANDLER_HTTP_TEXT: &str = "OK from synthetic HTTP upstream."; @@ -625,6 +626,86 @@ impl CopilotRequestHandler for RecordingHandler { } } +#[tokio::test] +async fn preserves_message_source_through_runtime_delivery() { + if super::support::skip_inprocess("LLM inference providers are process-global in-process") { + return; + } + with_e2e_context_no_snapshot(|ctx| { + Box::pin(async move { + let handler = Arc::new(RecordingHandler::default()); + let client = ctx.start_llm_client(handler.clone(), &[]).await; + + for typed_rpc in [false, true] { + for immediate in [false, true] { + for source in [None, Some("agent-sender-id")] { + let session = client + .create_session(ctx.approve_all_session_config()) + .await + .expect("create session"); + let events = session.subscribe(); + let message_id = if typed_rpc { + let mut request = SendRequest::default(); + request.prompt = "Say OK.".to_string(); + request.mode = Some(if immediate { + SendMode::Immediate + } else { + SendMode::Enqueue + }); + if let Some(source) = source { + request = request.with_source(source); + } + session + .rpc() + .send(request) + .await + .expect("RPC send") + .message_id + } else { + let mut options = say_ok().with_mode(if immediate { + DeliveryMode::Immediate + } else { + DeliveryMode::Enqueue + }); + if let Some(source) = source { + options = options.with_source(source); + } + session.send(options).await.expect("send") + }; + + let observed = collect_until_idle(events).await; + let messages: Vec<_> = observed + .iter() + .filter(|event| event.event_type == "user.message") + .map(|event| { + event.typed_data::().expect("user message") + }) + .collect(); + assert_eq!(messages.len(), 1, "expected one delivered user message"); + assert_eq!(messages[0].content, "Say OK."); + assert_eq!(messages[0].source.as_deref(), source); + assert_eq!(messages[0].message_id.as_deref(), Some(message_id.as_str())); + assert!( + observed.iter().any(|event| { + event.event_type == "assistant.message" + && event + .typed_data::() + .is_some_and(|data| data.content == SYNTHETIC_TEXT) + }), + "expected the synthetic provider response" + ); + + session.disconnect().await.expect("disconnect session"); + } + } + } + assert!(!handler.inference_records().is_empty()); + client.stop().await.expect("stop client"); + }) + }) + .await; +} + #[tokio::test] async fn threads_session_id_into_inference() { if super::support::skip_inprocess("LLM inference providers are process-global in-process") { diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index e7bafc683c..a094e076f5 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -19,7 +19,7 @@ use github_copilot_sdk::handler::{ }; use github_copilot_sdk::rpc::{ CanvasProviderInvokeActionRequest, CanvasProviderOpenRequest, CanvasProviderOpenResult, - OpenCanvasInstance, + OpenCanvasInstance, SendMode, SendRequest, }; use github_copilot_sdk::session_events::{ ManagedSettingsResolvedSource, McpOauthRequiredData, ReasoningSummary, SessionLimitsConfig, @@ -1865,6 +1865,139 @@ async fn send_injects_session_id() { timeout(TIMEOUT, handle).await.unwrap().unwrap().unwrap(); } +#[tokio::test] +async fn send_preserves_source_for_each_delivery_mode() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + + for (mode, wire_mode) in [ + (None, None), + (Some(DeliveryMode::Enqueue), Some("enqueue")), + (Some(DeliveryMode::Immediate), Some("immediate")), + ] { + for source in [None, Some("agent-sender-id")] { + let mut options = MessageOptions::new("Agent update"); + assert!(options.source.is_none()); + if let Some(mode) = mode { + options = options.with_mode(mode); + } + if let Some(source) = source { + options = options.with_source(source); + } + let handle = tokio::spawn({ + let session = session.clone(); + async move { session.send(options).await } + }); + + let request = timeout(TIMEOUT, server.read_request()).await.unwrap(); + let mut expected = serde_json::json!({ + "sessionId": server.session_id, + "prompt": "Agent update" + }); + if let Some(mode) = wire_mode { + expected["mode"] = serde_json::json!(mode); + } + if let Some(source) = source { + expected["source"] = serde_json::json!(source); + } + assert_eq!(request["method"], "session.send"); + assert_eq!(request["params"], expected); + server + .respond(&request, serde_json::json!({ "messageId": "message-1" })) + .await; + assert_eq!( + timeout(TIMEOUT, handle).await.unwrap().unwrap().unwrap(), + "message-1" + ); + } + } +} + +#[tokio::test] +async fn send_string_prompt_omits_source() { + let (session, mut server) = create_session_pair().await; + let handle = tokio::spawn(async move { session.send("Human message").await }); + + let request = timeout(TIMEOUT, server.read_request()).await.unwrap(); + assert_eq!(request["method"], "session.send"); + assert_eq!( + request["params"], + serde_json::json!({ + "sessionId": server.session_id, + "prompt": "Human message" + }) + ); + server + .respond(&request, serde_json::json!({ "messageId": "message-1" })) + .await; + assert_eq!( + timeout(TIMEOUT, handle).await.unwrap().unwrap().unwrap(), + "message-1" + ); +} + +#[tokio::test] +async fn rpc_send_preserves_source_and_attachments_for_each_delivery_mode() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let attachments = vec![serde_json::json!({ + "type": "extension_context", + "extensionId": "project:example", + "title": "Context", + "capturedAt": "2026-09-04T00:00:00Z", + "payload": { "content": "Agent context" } + })]; + + for (mode, wire_mode) in [ + (None, None), + (Some(SendMode::Enqueue), Some("enqueue")), + (Some(SendMode::Immediate), Some("immediate")), + ] { + for source in [None, Some("agent-sender-id")] { + let mut params = SendRequest::default(); + params.prompt = "Agent update".to_string(); + params.attachments = Some(attachments.clone()); + params.display_prompt = Some("Update from sender".to_string()); + params.mode = mode.clone(); + if let Some(source) = source { + params = params.with_source(source); + } + let handle = tokio::spawn({ + let session = session.clone(); + async move { session.rpc().send(params).await } + }); + + let request = timeout(TIMEOUT, server.read_request()).await.unwrap(); + let mut expected = serde_json::json!({ + "sessionId": server.session_id, + "prompt": "Agent update", + "displayPrompt": "Update from sender", + "attachments": attachments + }); + if let Some(mode) = wire_mode { + expected["mode"] = serde_json::json!(mode); + } + if let Some(source) = source { + expected["source"] = serde_json::json!(source); + } + assert_eq!(request["method"], "session.send"); + assert_eq!(request["params"], expected); + server + .respond(&request, serde_json::json!({ "messageId": "message-1" })) + .await; + assert_eq!( + timeout(TIMEOUT, handle) + .await + .unwrap() + .unwrap() + .unwrap() + .message_id, + "message-1" + ); + } + } +} + #[tokio::test] async fn send_serializes_request_headers() { use std::collections::HashMap; @@ -3463,6 +3596,46 @@ async fn send_and_wait_returns_last_assistant_message_on_idle() { assert_eq!(event.data["message"], "Hello back!"); } +#[tokio::test] +async fn send_and_wait_preserves_source_and_allows_idle_without_response() { + let (session, mut server) = create_session_pair().await; + let handle = tokio::spawn(async move { + session + .send_and_wait( + MessageOptions::new("Agent update") + .with_source("agent-sender-id") + .with_wait_timeout(TIMEOUT), + ) + .await + }); + + let request = timeout(TIMEOUT, server.read_request()).await.unwrap(); + assert_eq!(request["method"], "session.send"); + assert_eq!( + request["params"], + serde_json::json!({ + "sessionId": server.session_id, + "prompt": "Agent update", + "source": "agent-sender-id" + }) + ); + server + .respond(&request, serde_json::json!({ "messageId": "message-1" })) + .await; + server + .send_event("session.idle", serde_json::json!({})) + .await; + + assert!( + timeout(TIMEOUT, handle) + .await + .unwrap() + .unwrap() + .unwrap() + .is_none() + ); +} + #[tokio::test] async fn send_and_wait_returns_error_on_session_error() { let (session, mut server) = create_session_pair().await;