agent2agent 1.0.7 → 1.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. checksums.yaml +4 -4
  2. data/lib/a2a/agent.rb +165 -117
  3. data/lib/a2a/client.rb +470 -51
  4. data/lib/a2a/errors/json_rpc_error.rb +71 -0
  5. data/lib/a2a/errors/rest_error.rb +68 -0
  6. data/lib/a2a/errors.rb +535 -0
  7. data/lib/a2a/faraday/middleware/json_rpc/request.rb +96 -0
  8. data/lib/a2a/faraday/middleware/json_rpc/response.rb +131 -0
  9. data/lib/a2a/faraday/middleware/rest/request.rb +166 -0
  10. data/lib/a2a/faraday/middleware/rest/response.rb +144 -0
  11. data/lib/a2a/faraday/middleware/schema_request.rb +69 -0
  12. data/lib/a2a/middleware/extract_message.rb +120 -0
  13. data/lib/a2a/middleware/fetch_task.rb +228 -0
  14. data/lib/a2a/middleware/limit_history_length.rb +123 -0
  15. data/lib/a2a/middleware/limit_pagination_size.rb +133 -0
  16. data/lib/a2a/middleware/sse_stream.rb +235 -0
  17. data/lib/a2a/middleware.rb +7 -0
  18. data/lib/a2a/schema/definition.rb +35 -1
  19. data/lib/a2a/schema.rb +126 -0
  20. data/lib/a2a/{bindings → server/bindings}/json_rpc.rb +12 -8
  21. data/lib/a2a/{bindings → server/bindings}/rest.rb +12 -8
  22. data/lib/a2a/server/dispatcher.rb +52 -52
  23. data/lib/a2a/server/env.rb +4 -6
  24. data/lib/a2a/server/triage.rb +2 -1
  25. data/lib/a2a/server.rb +10 -10
  26. data/lib/a2a/sse/event_parser.rb +202 -0
  27. data/lib/a2a/sse/json_rpc_stream.rb +27 -5
  28. data/lib/a2a/sse/rest_stream.rb +17 -5
  29. data/lib/a2a/sse/stream.rb +135 -7
  30. data/lib/a2a/sse.rb +1 -0
  31. data/lib/a2a/test_helpers.rb +89 -0
  32. data/lib/a2a/version.rb +1 -1
  33. data/lib/a2a.rb +6 -2
  34. data/lib/traces/provider/a2a/{bindings → server/bindings}/json_rpc.rb +2 -2
  35. data/lib/traces/provider/a2a/{bindings → server/bindings}/rest.rb +2 -2
  36. data/lib/traces/provider/a2a.rb +2 -2
  37. metadata +49 -22
  38. data/lib/a2a/server/cancel_task.rb +0 -14
  39. data/lib/a2a/server/create_task_push_notification_config.rb +0 -14
  40. data/lib/a2a/server/delete_task_push_notification_config.rb +0 -14
  41. data/lib/a2a/server/get_extended_agent_card.rb +0 -15
  42. data/lib/a2a/server/get_task.rb +0 -14
  43. data/lib/a2a/server/get_task_push_notification_config.rb +0 -14
  44. data/lib/a2a/server/list_task_push_notification_configs.rb +0 -14
  45. data/lib/a2a/server/list_tasks.rb +0 -14
  46. data/lib/a2a/server/send_message.rb +0 -14
  47. data/lib/a2a/server/send_streaming_message.rb +0 -14
  48. data/lib/a2a/server/subscribe_to_task.rb +0 -14
  49. data/lib/a2a/store/processor.rb +0 -136
  50. data/lib/a2a/store/pub_sub.rb +0 -149
  51. data/lib/a2a/store/sqlite.rb +0 -533
  52. data/lib/a2a/store/webhooks.rb +0 -94
  53. data/lib/a2a/store.rb +0 -6
  54. data/lib/a2a/task_store.rb +0 -315
@@ -6,115 +6,115 @@ require "a2a"
6
6
 
7
7
  module A2A
8
8
  class Server
9
- # Routes incoming A2A operations to registered handler objects.
9
+ # Routes incoming A2A operations to their registered handler stacks.
10
10
  #
11
- # Each handler declares the operations it handles via `#operations`.
12
- # When an operation arrives, the dispatcher finds all matching handlers
13
- # and calls them. Errors in one handler do not prevent others from running.
11
+ # Each operation maps to exactly one Rack app (a compiled middleware
12
+ # stack built by the Agent DSL). The Dispatcher is a terminal Rack app
13
+ # that reads env["a2a.operation"] set by Triage, looks up the matching
14
+ # route, and calls it.
14
15
  #
15
- # The Dispatcher is a Rack app (terminal, not middleware). It reads
16
- # env["a2a.operation"] set by Triage and fans out to matching handlers.
16
+ # Returns domain objects (A2A::Schema::Definition or A2A::Error)
17
+ # the binding layer is responsible for formatting these into HTTP.
17
18
  #
18
19
  class Dispatcher
19
20
  def initialize
20
- @handlers = Hash.new { |h, k| h[k] = [] }
21
+ @routes = {}
21
22
  end
22
23
 
23
24
  # Register a handler object.
24
25
  #
25
26
  # The handler must respond to:
26
27
  # #operations -> Array(String) (e.g. ["SendMessage", "GetTask"])
27
- # #call(env) -> void (sets env["a2a.result"])
28
+ # #call(env) -> A2A::Schema::Definition | A2A::Error
28
29
  #
29
30
  def register(handler)
30
31
  handler.operations.each do |op|
31
- @handlers[op] << handler
32
+ @routes[op] = handler
32
33
  Console.info(self) { "Registered #{handler.class.name} for #{op}" }
33
34
  end
34
35
  end
35
36
 
36
37
  def call(env)
37
38
  operation = env["a2a.operation"]
39
+ app = @routes[operation]
38
40
 
39
- if operation
40
- dispatch(operation, env)
41
+ unless app
42
+ raise A2A::UnsupportedOperationError.new(
43
+ message: "Operation not supported: #{operation}"
44
+ )
41
45
  end
42
46
 
43
- [200, {}, []]
47
+ app.call(env)
48
+ rescue A2A::Error => e
49
+ e
50
+ rescue => e
51
+ Console.error(self, "Handler raised #{e.class}: #{e.message}", e)
52
+ A2A::Error.new("Internal error", code: -32603, http_status: 500)
44
53
  end
45
54
 
46
55
  def handler_count
47
- @handlers.values.flatten.size
56
+ @routes.size
48
57
  end
49
-
50
- private
51
-
52
- def dispatch(operation, env)
53
- handlers = @handlers[operation]
54
-
55
- if handlers.empty?
56
- Console.debug(self) { "No handler for operation: #{operation}" }
57
- else
58
- handlers.each do |handler|
59
- begin
60
- handler.call(env)
61
- rescue => e
62
- Console.error(self, "Handler #{handler.class.name} raised #{e.class}", e)
63
- end
64
- end
65
- end
66
- end
67
58
  end
68
59
  end
69
60
  end
70
61
 
71
62
  test do
72
63
  describe "A2A::Server::Dispatcher" do
73
- it "registers and dispatches to handlers" do
74
- received = []
64
+ it "registers and dispatches to a handler" do
75
65
  handler = Object.new
76
66
  handler.define_singleton_method(:operations) { ["SendMessage"] }
77
- handler.define_singleton_method(:call) { |env| received << env }
67
+ handler.define_singleton_method(:call) { |env| :dispatched }
78
68
 
79
69
  dispatcher = A2A::Server::Dispatcher.new
80
70
  dispatcher.register(handler)
81
71
  dispatcher.handler_count.should == 1
82
72
 
83
73
  env = { "a2a.operation" => "SendMessage" }
84
- dispatcher.call(env)
85
- received.length.should == 1
74
+ result = dispatcher.call(env)
75
+ result.should == :dispatched
86
76
  end
87
77
 
88
- it "ignores operations with no matching handler" do
78
+ it "returns UnsupportedOperationError for unknown operations" do
89
79
  dispatcher = A2A::Server::Dispatcher.new
90
80
  env = { "a2a.operation" => "UnknownOp" }
91
- lambda { dispatcher.call(env) }.should.not.raise
81
+ result = dispatcher.call(env)
82
+ result.should.be.is_a(A2A::UnsupportedOperationError)
83
+ result.code.should == -32004
92
84
  end
93
85
 
94
- it "continues dispatching when a handler raises" do
95
- results = []
96
- bad_handler = Object.new
97
- bad_handler.define_singleton_method(:operations) { ["SendMessage"] }
98
- bad_handler.define_singleton_method(:call) { |_| raise "boom" }
86
+ it "returns A2A::Error when handler raises one" do
87
+ handler = Object.new
88
+ handler.define_singleton_method(:operations) { ["SendMessage"] }
89
+ handler.define_singleton_method(:call) { |_| raise A2A::TaskNotFoundError.new("t-1") }
90
+
91
+ dispatcher = A2A::Server::Dispatcher.new
92
+ dispatcher.register(handler)
93
+
94
+ result = dispatcher.call({ "a2a.operation" => "SendMessage" })
95
+ result.should.be.is_a(A2A::TaskNotFoundError)
96
+ result.code.should == -32001
97
+ end
99
98
 
100
- good_handler = Object.new
101
- good_handler.define_singleton_method(:operations) { ["SendMessage"] }
102
- good_handler.define_singleton_method(:call) { |e| results << e }
99
+ it "returns internal error when handler raises unexpected exception" do
100
+ handler = Object.new
101
+ handler.define_singleton_method(:operations) { ["SendMessage"] }
102
+ handler.define_singleton_method(:call) { |_| raise "boom" }
103
103
 
104
104
  dispatcher = A2A::Server::Dispatcher.new
105
- dispatcher.register(bad_handler)
106
- dispatcher.register(good_handler)
105
+ dispatcher.register(handler)
107
106
 
108
- env = { "a2a.operation" => "SendMessage" }
109
- dispatcher.call(env)
110
- results.length.should == 1
107
+ result = dispatcher.call({ "a2a.operation" => "SendMessage" })
108
+ result.should.be.is_a(A2A::Error)
109
+ result.code.should == -32603
110
+ result.http_status.should == 500
111
111
  end
112
112
 
113
113
  it "dispatches to multiple operations from one handler" do
114
114
  received = []
115
115
  handler = Object.new
116
116
  handler.define_singleton_method(:operations) { ["SendMessage", "GetTask"] }
117
- handler.define_singleton_method(:call) { |env| received << env["a2a.operation"] }
117
+ handler.define_singleton_method(:call) { |env| received << env["a2a.operation"]; :ok }
118
118
 
119
119
  dispatcher = A2A::Server::Dispatcher.new
120
120
  dispatcher.register(handler)
@@ -7,19 +7,17 @@ module A2A
7
7
  class Server
8
8
  # Rack middleware that injects shared A2A context into the env.
9
9
  #
10
- # Sets env["a2a.store"] and env["a2a.agent_card"] so downstream
11
- # middleware and handlers can access them without coupling to
12
- # any particular configuration mechanism.
10
+ # Sets env["a2a.agent_card"] so downstream middleware and handlers
11
+ # can access it without coupling to any particular configuration
12
+ # mechanism.
13
13
  #
14
14
  class Env
15
- def initialize(app, agent_card: {}, store: TaskStore.new)
15
+ def initialize(app, agent_card: {})
16
16
  @app = app
17
17
  @agent_card = agent_card
18
- @store = store
19
18
  end
20
19
 
21
20
  def call(env)
22
- env["a2a.store"] = @store
23
21
  env["a2a.agent_card"] = @agent_card
24
22
  @app.call(env)
25
23
  end
@@ -29,10 +29,11 @@ module A2A
29
29
  end
30
30
 
31
31
  def call(env)
32
+ Console.info(self) { "Triaging: #{env}" }
32
33
  op = resolve_operation(env)
33
34
 
34
35
  unless op
35
- return [404, { "content-type" => "text/plain" }, ["Unknown operation"]]
36
+ return A2A::UnsupportedOperationError.new(message: "Unknown operation")
36
37
  end
37
38
 
38
39
  env["a2a.operation"] = op.name
data/lib/a2a/server.rb CHANGED
@@ -24,8 +24,10 @@ module A2A
24
24
  # Usage:
25
25
  #
26
26
  # agent = A2A::Agent.new do
27
- # on "SendMessage" do |request|
28
- # respond A2A::Schema["Send Message Response"].new({})
27
+ # on "SendMessage" do
28
+ # respond_with -> (env) {
29
+ # A2A::Schema["Send Message Response"].new({})
30
+ # }
29
31
  # end
30
32
  # end
31
33
  #
@@ -35,9 +37,8 @@ module A2A
35
37
  # run app
36
38
  #
37
39
  class Server
38
- def initialize(agent_card: {}, store: TaskStore.new)
40
+ def initialize(agent_card: {})
39
41
  @agent_card = agent_card
40
- @store = store
41
42
  @dispatcher = Dispatcher.new
42
43
  @app = build_app
43
44
  end
@@ -63,15 +64,14 @@ module A2A
63
64
  private
64
65
 
65
66
  def build_app
66
- require "a2a/bindings/json_rpc"
67
- require "a2a/bindings/rest"
67
+ require "a2a/server/bindings/json_rpc"
68
+ require "a2a/server/bindings/rest"
68
69
 
69
70
  agent_card = @agent_card
70
- store = @store
71
71
  dispatcher = @dispatcher
72
72
 
73
73
  Rack::Builder.app do
74
- use A2A::Server::Env, agent_card: agent_card, store: store
74
+ use A2A::Server::Env, agent_card: agent_card
75
75
 
76
76
  map "/.well-known/agent-card.json" do
77
77
  run ->(env) {
@@ -82,13 +82,13 @@ module A2A
82
82
  end
83
83
 
84
84
  map "/a2a" do
85
- use A2A::Bindings::JsonRpc
85
+ use A2A::Server::Bindings::JsonRpc
86
86
  use A2A::Server::Triage
87
87
  run dispatcher
88
88
  end
89
89
 
90
90
  map "/" do
91
- use A2A::Bindings::Rest
91
+ use A2A::Server::Bindings::Rest
92
92
  use A2A::Server::Triage
93
93
  run dispatcher
94
94
  end
@@ -0,0 +1,202 @@
1
+ # frozen_string_literal: true
2
+
3
+ require "json"
4
+
5
+ module A2A
6
+ module SSE
7
+ # Parses raw SSE text chunks into typed StreamResponse schema objects.
8
+ #
9
+ # Faraday's `on_data` callback delivers raw text in arbitrary-sized
10
+ # chunks. This parser buffers them, extracts complete SSE events
11
+ # (delimited by blank lines), parses the JSON from `data:` lines,
12
+ # and yields A2A::Schema["Stream Response"] instances.
13
+ #
14
+ # For the JSON-RPC binding, the parser unwraps the JSON-RPC 2.0
15
+ # envelope (`{"jsonrpc":"2.0","id":N,"result":{...}}`) to extract
16
+ # the StreamResponse payload from `result`.
17
+ #
18
+ # Usage:
19
+ #
20
+ # parser = A2A::SSE::EventParser.new(binding: :rest)
21
+ #
22
+ # parser.feed("data: {\"task\":{\"id\":\"t1\"}}\n\n") do |event|
23
+ # event.task.id #=> "t1"
24
+ # end
25
+ #
26
+ class EventParser
27
+ def initialize(binding:)
28
+ @binding = binding
29
+ @buffer = String.new
30
+ end
31
+
32
+ # Feed a raw chunk of SSE text. Yields a StreamResponse for each
33
+ # complete event found in the buffer.
34
+ #
35
+ # @param chunk [String] raw SSE text from the wire
36
+ # @yieldparam event [A2A::Schema::Definition] a Stream Response instance
37
+ #
38
+ def feed(chunk)
39
+ @buffer << chunk
40
+
41
+ # SSE events are terminated by a blank line (\n\n).
42
+ # We split on that boundary, keeping any incomplete trailing data
43
+ # in the buffer for the next call.
44
+ while (idx = @buffer.index("\n\n"))
45
+ raw_event = @buffer.slice!(0, idx + 2)
46
+ payload = parse_event(raw_event)
47
+ yield A2A::Schema["Stream Response"].new(payload) if payload
48
+ end
49
+ end
50
+
51
+ private
52
+
53
+ # Extract JSON from `data:` lines and optionally unwrap JSON-RPC.
54
+ def parse_event(raw)
55
+ data_lines = raw.each_line
56
+ .select { |line| line.start_with?("data:") }
57
+ .map { |line| line.sub(/\Adata:\s?/, "").chomp }
58
+
59
+ return nil if data_lines.empty?
60
+
61
+ parsed = JSON.parse(data_lines.join("\n"))
62
+
63
+ case @binding
64
+ when :json_rpc then parsed["result"]
65
+ else parsed
66
+ end
67
+ rescue JSON::ParserError
68
+ nil
69
+ end
70
+ end
71
+ end
72
+ end
73
+
74
+ test do
75
+ describe "A2A::SSE::EventParser" do
76
+ # --- REST binding ---
77
+
78
+ it "parses a single complete SSE event" do
79
+ parser = A2A::SSE::EventParser.new(binding: :rest)
80
+ events = []
81
+
82
+ parser.feed("data: {\"task\":{\"id\":\"t1\",\"contextId\":\"c1\"}}\n\n") do |event|
83
+ events << event
84
+ end
85
+
86
+ events.size.should == 1
87
+ events[0].should.be.kind_of(A2A::Schema::Definition)
88
+ events[0].task.should.not.be.nil
89
+ events[0].task.id.should == "t1"
90
+ end
91
+
92
+ it "parses multiple events in a single chunk" do
93
+ parser = A2A::SSE::EventParser.new(binding: :rest)
94
+ events = []
95
+
96
+ chunk = [
97
+ "data: {\"task\":{\"id\":\"t1\",\"contextId\":\"c1\"}}\n\n",
98
+ "data: {\"statusUpdate\":{\"taskId\":\"t1\",\"contextId\":\"c1\",\"status\":{\"state\":\"TASK_STATE_COMPLETED\"}}}\n\n",
99
+ ].join
100
+
101
+ parser.feed(chunk) { |e| events << e }
102
+
103
+ events.size.should == 2
104
+ events[0].task.should.not.be.nil
105
+ events[1].status_update.should.not.be.nil
106
+ end
107
+
108
+ it "handles partial chunks via buffering" do
109
+ parser = A2A::SSE::EventParser.new(binding: :rest)
110
+ events = []
111
+
112
+ # First chunk: incomplete event
113
+ parser.feed("data: {\"task\":{\"id\"") { |e| events << e }
114
+ events.size.should == 0
115
+
116
+ # Second chunk: completes the event
117
+ parser.feed(":\"t1\",\"contextId\":\"c1\"}}\n\n") { |e| events << e }
118
+ events.size.should == 1
119
+ events[0].task.id.should == "t1"
120
+ end
121
+
122
+ it "handles events with event: and id: fields" do
123
+ parser = A2A::SSE::EventParser.new(binding: :rest)
124
+ events = []
125
+
126
+ chunk = "event: message\nid: 42\ndata: {\"task\":{\"id\":\"t1\",\"contextId\":\"c1\"}}\n\n"
127
+ parser.feed(chunk) { |e| events << e }
128
+
129
+ events.size.should == 1
130
+ events[0].task.id.should == "t1"
131
+ end
132
+
133
+ it "skips events with no data lines" do
134
+ parser = A2A::SSE::EventParser.new(binding: :rest)
135
+ events = []
136
+
137
+ parser.feed(": comment\n\n") { |e| events << e }
138
+ events.size.should == 0
139
+ end
140
+
141
+ it "skips events with invalid JSON" do
142
+ parser = A2A::SSE::EventParser.new(binding: :rest)
143
+ events = []
144
+
145
+ parser.feed("data: not-json\n\n") { |e| events << e }
146
+ events.size.should == 0
147
+ end
148
+
149
+ # --- JSON-RPC binding ---
150
+
151
+ it "unwraps JSON-RPC 2.0 envelope" do
152
+ parser = A2A::SSE::EventParser.new(binding: :json_rpc)
153
+ events = []
154
+
155
+ envelope = {
156
+ "jsonrpc" => "2.0",
157
+ "id" => 1,
158
+ "result" => {
159
+ "task" => { "id" => "t1", "contextId" => "c1" },
160
+ },
161
+ }
162
+
163
+ parser.feed("data: #{JSON.generate(envelope)}\n\n") { |e| events << e }
164
+
165
+ events.size.should == 1
166
+ events[0].task.should.not.be.nil
167
+ events[0].task.id.should == "t1"
168
+ end
169
+
170
+ it "handles multiple JSON-RPC events" do
171
+ parser = A2A::SSE::EventParser.new(binding: :json_rpc)
172
+ events = []
173
+
174
+ [
175
+ { "jsonrpc" => "2.0", "id" => 1, "result" => { "task" => { "id" => "t1", "contextId" => "c1" } } },
176
+ { "jsonrpc" => "2.0", "id" => 1, "result" => { "statusUpdate" => { "taskId" => "t1", "contextId" => "c1", "status" => { "state" => "TASK_STATE_COMPLETED" } } } },
177
+ ].each do |envelope|
178
+ parser.feed("data: #{JSON.generate(envelope)}\n\n") { |e| events << e }
179
+ end
180
+
181
+ events.size.should == 2
182
+ events[0].task.should.not.be.nil
183
+ events[1].status_update.should.not.be.nil
184
+ end
185
+
186
+ # --- Multi-line data ---
187
+
188
+ it "handles multi-line data fields" do
189
+ parser = A2A::SSE::EventParser.new(binding: :rest)
190
+ events = []
191
+
192
+ # JSON split across multiple data: lines
193
+ chunk = "data: {\"task\":\n" \
194
+ "data: {\"id\":\"t1\",\"contextId\":\"c1\"}}\n\n"
195
+
196
+ parser.feed(chunk) { |e| events << e }
197
+
198
+ events.size.should == 1
199
+ events[0].task.id.should == "t1"
200
+ end
201
+ end
202
+ end
@@ -11,10 +11,12 @@ module A2A
11
11
  #
12
12
  # Usage:
13
13
  #
14
- # stream = A2A::SSE::JsonRpcStream.new(json_rpc_id: 1)
14
+ # stream = A2A::SSE::JsonRpcStream.new(
15
+ # task_id: "t1", context_id: "c1", json_rpc_id: 1
16
+ # )
15
17
  #
16
18
  # Async do
17
- # stream.event({ "task" => { ... } })
19
+ # stream.task(status: { state: "TASK_STATE_WORKING", timestamp: "..." })
18
20
  # stream.finish
19
21
  # end
20
22
  #
@@ -45,9 +47,11 @@ end
45
47
  test do
46
48
  describe "A2A::SSE::JsonRpcStream" do
47
49
  it "wraps events in JSON-RPC 2.0 envelopes" do
48
- stream = A2A::SSE::JsonRpcStream.new(json_rpc_id: 42)
50
+ stream = A2A::SSE::JsonRpcStream.new(
51
+ task_id: "t1", context_id: "c1", json_rpc_id: 42
52
+ )
49
53
 
50
- stream.event({ "task" => { "id" => "t1" } })
54
+ stream.task(status: { state: "TASK_STATE_WORKING", timestamp: "2025-01-01T00:00:00Z" })
51
55
  stream.finish
52
56
 
53
57
  chunk = stream.read
@@ -58,9 +62,27 @@ test do
58
62
  end
59
63
 
60
64
  it "is a subclass of SSE::Stream" do
61
- stream = A2A::SSE::JsonRpcStream.new(json_rpc_id: 1)
65
+ stream = A2A::SSE::JsonRpcStream.new(
66
+ task_id: "t1", context_id: "c1", json_rpc_id: 1
67
+ )
62
68
  stream.is_a?(A2A::SSE::Stream).should == true
63
69
  stream.is_a?(Protocol::HTTP::Body::Readable).should == true
64
70
  end
71
+
72
+ it "typed methods chain through JSON-RPC envelope" do
73
+ stream = A2A::SSE::JsonRpcStream.new(
74
+ task_id: "t1", context_id: "c1", json_rpc_id: 7
75
+ )
76
+
77
+ stream.status_update(status: { state: "TASK_STATE_COMPLETED", timestamp: "2025-01-01T00:00:00Z" })
78
+ stream.finish
79
+
80
+ chunk = stream.read
81
+ parsed = JSON.parse(chunk.sub(/\Adata: /, "").strip)
82
+ parsed["jsonrpc"].should == "2.0"
83
+ parsed["id"].should == 7
84
+ parsed["result"]["statusUpdate"]["taskId"].should == "t1"
85
+ parsed["result"]["statusUpdate"]["contextId"].should == "c1"
86
+ end
65
87
  end
66
88
  end
@@ -15,10 +15,10 @@ module A2A
15
15
  #
16
16
  # Usage:
17
17
  #
18
- # stream = A2A::SSE::RestStream.new
18
+ # stream = A2A::SSE::RestStream.new(task_id: "t1", context_id: "c1")
19
19
  #
20
20
  # Async do
21
- # stream.event({ "task" => { ... } })
21
+ # stream.task(status: { state: "TASK_STATE_WORKING", timestamp: "..." })
22
22
  # stream.finish
23
23
  # end
24
24
  #
@@ -32,9 +32,9 @@ end
32
32
  test do
33
33
  describe "A2A::SSE::RestStream" do
34
34
  it "emits bare JSON events (no envelope)" do
35
- stream = A2A::SSE::RestStream.new
35
+ stream = A2A::SSE::RestStream.new(task_id: "t1", context_id: "c1")
36
36
 
37
- stream.event({ "statusUpdate" => { "taskId" => "t1" } })
37
+ stream.status_update(status: { state: "TASK_STATE_COMPLETED", timestamp: "2025-01-01T00:00:00Z" })
38
38
  stream.finish
39
39
 
40
40
  chunk = stream.read
@@ -43,8 +43,20 @@ test do
43
43
  end
44
44
 
45
45
  it "is a subclass of SSE::Stream" do
46
- stream = A2A::SSE::RestStream.new
46
+ stream = A2A::SSE::RestStream.new(task_id: "t1", context_id: "c1")
47
47
  stream.is_a?(A2A::SSE::Stream).should == true
48
48
  end
49
+
50
+ it "injects task_id and context_id into typed events" do
51
+ stream = A2A::SSE::RestStream.new(task_id: "t1", context_id: "c1")
52
+
53
+ stream.task(status: { state: "TASK_STATE_WORKING" })
54
+ stream.finish
55
+
56
+ chunk = stream.read
57
+ parsed = JSON.parse(chunk.sub(/\Adata: /, "").strip)
58
+ parsed["task"]["id"].should == "t1"
59
+ parsed["task"]["contextId"].should == "c1"
60
+ end
49
61
  end
50
62
  end