Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 09ed60e

Browse files
jujnjujn
authored
fix(A2A): fix A2aAgent lifecycle events to support TTSHook and reasoning completion (agentscope-ai#1150)
## Description Close agentscope-ai#1120 This PR fixes the issue where plugins like TTSHook failed in A2aAgent by introducing a state machine in ClientEventContext to guarantee the execution order of PreReasoningEvent, ReasoningChunkEvent, and PostReasoningEvent. ## Checklist Please check the following items before code is ready to be reviewed. - [ ] Code has been formatted with `mvn spotless:apply` - [ ] All tests are passing (`mvn test`) - [ ] Javadoc comments are complete and follow project conventions - [ ] Related documentation has been updated (e.g. links, examples, etc.) - [ ] Code is ready for review
1 parent 0058bff commit 09ed60e

6 files changed

Lines changed: 212 additions & 10 deletions

File tree

‎agentscope-extensions/agentscope-extensions-a2a/agentscope-extensions-a2a-client/src/main/java/io/agentscope/core/a2a/agent/A2aAgent.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ protected Mono<Msg> doCall(List<Msg> msgs) {
120120
LoggerUtil.debug(log, "[{}] A2aAgent call with input messages: ", currentRequestId);
121121
LoggerUtil.logTextMsgDetail(log, memory.getMessages());
122122
clientEventContext.setHooks(getSortedHooks());
123+
clientEventContext.setInputMessages(memory.getMessages());
123124
return Mono.defer(
124125
() -> {
125126
Message message =

‎agentscope-extensions/agentscope-extensions-a2a/agentscope-extensions-a2a-client/src/main/java/io/agentscope/core/a2a/agent/event/ClientEventContext.java‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,13 @@
1919
import io.a2a.spec.Task;
2020
import io.agentscope.core.a2a.agent.A2aAgent;
2121
import io.agentscope.core.hook.Hook;
22+
import io.agentscope.core.hook.PostReasoningEvent;
23+
import io.agentscope.core.hook.PreReasoningEvent;
24+
import io.agentscope.core.hook.ReasoningChunkEvent;
2225
import io.agentscope.core.message.Msg;
2326
import java.util.List;
27+
import java.util.concurrent.atomic.AtomicBoolean;
28+
import reactor.core.publisher.Mono;
2429
import reactor.core.publisher.MonoSink;
2530

2631
/**
@@ -40,6 +45,16 @@ public class ClientEventContext {
4045

4146
private Task task;
4247

48+
/**
49+
* Temporarily store the complete historical dialogue context at the time of this call,
50+
* specifically for use in constructing PreReasoning Events using the {@link #publishPreReasoning()} method.
51+
*/
52+
private List<Msg> inputMessages;
53+
54+
// Ensure that lifecycle events are triggered only once
55+
private final AtomicBoolean preReasoningFired = new AtomicBoolean(false);
56+
private final AtomicBoolean postReasoningFired = new AtomicBoolean(false);
57+
4358
public ClientEventContext(String currentRequestId, A2aAgent agent) {
4459
this.currentRequestId = currentRequestId;
4560
this.agent = agent;
@@ -76,4 +91,72 @@ public Task getTask() {
7691
public void setTask(Task task) {
7792
this.task = task;
7893
}
94+
95+
public void setInputMessages(List<Msg> inputMessages) {
96+
this.inputMessages = inputMessages;
97+
}
98+
99+
// ==========================================
100+
// Unified Event Publishing API
101+
// ==========================================
102+
103+
/**
104+
* Trigger PreReasoningEvent (triggered only once)
105+
*/
106+
void publishPreReasoning() {
107+
if (hooks != null && !hooks.isEmpty() && preReasoningFired.compareAndSet(false, true)) {
108+
List<Msg> msgs = inputMessages == null ? List.of() : inputMessages;
109+
PreReasoningEvent preEvent = new PreReasoningEvent(agent, "A2A", null, msgs);
110+
111+
Mono<PreReasoningEvent> eventMono = Mono.just(preEvent);
112+
for (Hook hook : hooks) {
113+
eventMono = eventMono.flatMap(hook::onEvent);
114+
}
115+
eventMono.block();
116+
}
117+
}
118+
119+
/**
120+
* Trigger ReasoningChunkEvent (streaming process)
121+
*/
122+
void publishReasoningChunk(Msg chunkMsg) {
123+
if (hooks != null && !hooks.isEmpty()) {
124+
publishPreReasoning(); // If not sent Pre before, send Pre first
125+
ReasoningChunkEvent chunkEvent =
126+
new ReasoningChunkEvent(agent, "A2A", null, chunkMsg, chunkMsg);
127+
128+
Mono<ReasoningChunkEvent> eventMono = Mono.just(chunkEvent);
129+
for (Hook hook : hooks) {
130+
eventMono = eventMono.flatMap(hook::onEvent);
131+
}
132+
eventMono.block();
133+
}
134+
}
135+
136+
/**
137+
* Trigger PostReasoningEvent (triggered only once) and return the final reasoning message
138+
* after hooks have had a chance to modify it.
139+
*
140+
* @param finalMsg the original final reasoning message
141+
* @return the hook-modified reasoning message, or {@code finalMsg} if no hooks ran or no
142+
* modification was applied
143+
*/
144+
Msg publishPostReasoning(Msg finalMsg) {
145+
if (hooks != null && !hooks.isEmpty() && postReasoningFired.compareAndSet(false, true)) {
146+
publishPreReasoning();
147+
PostReasoningEvent postEvent = new PostReasoningEvent(agent, "A2A", null, finalMsg);
148+
149+
Mono<PostReasoningEvent> eventMono = Mono.just(postEvent);
150+
for (Hook hook : hooks) {
151+
eventMono = eventMono.flatMap(hook::onEvent);
152+
}
153+
154+
postEvent = eventMono.block();
155+
if (postEvent != null) {
156+
Msg modifiedMsg = postEvent.getReasoningMessage();
157+
return modifiedMsg != null ? modifiedMsg : finalMsg;
158+
}
159+
}
160+
return finalMsg;
161+
}
79162
}

‎agentscope-extensions/agentscope-extensions-a2a/agentscope-extensions-a2a-client/src/main/java/io/agentscope/core/a2a/agent/event/MessageEventHandler.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ public void handle(MessageEvent event, ClientEventContext context) {
4242
Msg msg =
4343
MessageConvertUtil.convertFromMessage(
4444
event.getMessage(), context.getAgent().getName());
45+
46+
// Automatically trigger PreReasoningEvent and PostReasoningEvent
47+
msg = context.publishPostReasoning(msg);
48+
4549
context.getSink().success(msg);
4650
LoggerUtil.info(log, "[{}] A2aAgent complete call.", currentRequestId);
4751
LoggerUtil.debug(log, "[{}] A2aAgent complete with artifact messages: ", currentRequestId);

‎agentscope-extensions/agentscope-extensions-a2a/agentscope-extensions-a2a-client/src/main/java/io/agentscope/core/a2a/agent/event/TaskEventHandler.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,5 +44,7 @@ public void handle(TaskEvent event, ClientEventContext context) {
4444
context.getCurrentRequestId(),
4545
task.getId(),
4646
task.getStatus());
47+
48+
context.publishPreReasoning();
4749
}
4850
}

‎agentscope-extensions/agentscope-extensions-a2a/agentscope-extensions-a2a-client/src/main/java/io/agentscope/core/a2a/agent/event/TaskUpdateEventHandler.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import io.a2a.spec.UpdateEvent;
2424
import io.agentscope.core.a2a.agent.utils.LoggerUtil;
2525
import io.agentscope.core.a2a.agent.utils.MessageConvertUtil;
26-
import io.agentscope.core.hook.ReasoningChunkEvent;
2726
import io.agentscope.core.message.Msg;
2827
import java.util.HashMap;
2928
import java.util.List;
@@ -93,6 +92,9 @@ public void handle(TaskStatusUpdateEvent event, ClientEventContext context) {
9392
Msg msg =
9493
MessageConvertUtil.convertFromArtifact(
9594
context.getTask().getArtifacts(), context.getAgent().getName());
95+
96+
msg = context.publishPostReasoning(msg);
97+
9698
context.getSink().success(msg);
9799
LoggerUtil.info(log, "[{}] A2aAgent complete call.", currentRequestId);
98100
LoggerUtil.debug(
@@ -114,9 +116,8 @@ public void handle(TaskStatusUpdateEvent event, ClientEventContext context) {
114116
LoggerUtil.debug(
115117
log, "[{}] A2aAgent task status updated with messages: ", currentRequestId);
116118
LoggerUtil.logTextMsgDetail(log, List.of(msg));
117-
ReasoningChunkEvent chunkEvent =
118-
new ReasoningChunkEvent(context.getAgent(), "A2A", null, msg, msg);
119-
context.getHooks().forEach(hook -> hook.onEvent(chunkEvent).block());
119+
120+
context.publishReasoningChunk(msg);
120121
}
121122
}
122123
}
@@ -136,9 +137,8 @@ public void handle(TaskArtifactUpdateEvent event, ClientEventContext context) {
136137
LoggerUtil.debug(
137138
log, "[{}] A2aAgent artifact append with messages: ", currentRequestTaskId);
138139
LoggerUtil.logTextMsgDetail(log, List.of(msg));
139-
ReasoningChunkEvent chunkEvent =
140-
new ReasoningChunkEvent(context.getAgent(), "A2A", null, msg, msg);
141-
context.getHooks().forEach(hook -> hook.onEvent(chunkEvent).block());
140+
141+
context.publishReasoningChunk(msg);
142142
}
143143
}
144144
}

‎agentscope-extensions/agentscope-extensions-a2a/agentscope-extensions-a2a-client/src/test/java/io/agentscope/core/a2a/agent/A2aAgentTest.java‎

Lines changed: 115 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,10 @@
5858
import io.agentscope.core.agent.Event;
5959
import io.agentscope.core.hook.Hook;
6060
import io.agentscope.core.hook.HookEvent;
61+
import io.agentscope.core.hook.PostReasoningEvent;
6162
import io.agentscope.core.hook.PreCallEvent;
63+
import io.agentscope.core.hook.PreReasoningEvent;
64+
import io.agentscope.core.hook.ReasoningChunkEvent;
6265
import io.agentscope.core.message.Msg;
6366
import java.lang.reflect.Field;
6467
import java.util.HashMap;
@@ -69,6 +72,7 @@
6972
import java.util.concurrent.TimeUnit;
7073
import java.util.concurrent.TimeoutException;
7174
import java.util.concurrent.atomic.AtomicBoolean;
75+
import java.util.concurrent.atomic.AtomicInteger;
7276
import java.util.function.BiConsumer;
7377
import org.junit.jupiter.api.BeforeEach;
7478
import org.junit.jupiter.api.DisplayName;
@@ -282,9 +286,10 @@ void testStreamAgentWithDefaultTransport() {
282286
List<Event> streamResults =
283287
agent.stream(Msg.builder().textContent("test").build()).collectList().block();
284288
assertNotNull(streamResults);
285-
assertEquals(2, streamResults.size());
286-
assertFalse(streamResults.get(0).isLast());
287-
assertTrue(streamResults.get(1).isLast());
289+
assertEquals(3, streamResults.size());
290+
assertFalse(streamResults.get(0).isLast()); // ReasoningChunkEvent
291+
assertTrue(streamResults.get(1).isLast()); // PostReasoningEvent
292+
assertTrue(streamResults.get(2).isLast()); // AGENT_RESULT
288293
}
289294

290295
@Test
@@ -428,6 +433,113 @@ void testCallAgentWithDefaultTransportByObserve() {
428433
assertEquals(3, agent.getMemory().getMessages().size());
429434
}
430435

436+
@Test
437+
@DisplayName("Should trigger Pre, Chunk, Post reasoning events")
438+
void testAgentLifecycleHooksTriggeredCorrectly() {
439+
AtomicInteger preCount = new AtomicInteger(0);
440+
AtomicInteger chunkCount = new AtomicInteger(0);
441+
AtomicInteger postCount = new AtomicInteger(0);
442+
443+
Hook lifecycleMonitorHook =
444+
new Hook() {
445+
@Override
446+
public <T extends HookEvent> Mono<T> onEvent(T event) {
447+
if (event instanceof PreReasoningEvent) {
448+
preCount.incrementAndGet();
449+
} else if (event instanceof ReasoningChunkEvent) {
450+
chunkCount.incrementAndGet();
451+
} else if (event instanceof PostReasoningEvent) {
452+
postCount.incrementAndGet();
453+
}
454+
return Mono.just(event);
455+
}
456+
457+
@Override
458+
public int priority() {
459+
return 1;
460+
}
461+
};
462+
463+
A2aAgent agent =
464+
A2aAgent.builder()
465+
.name("test-lifecycle-agent")
466+
.agentCard(agentCard)
467+
.hook(new ReplaceA2aClientHook())
468+
.hook(lifecycleMonitorHook)
469+
.build();
470+
471+
Answer<Void> mockTaskResponse =
472+
invocation -> {
473+
@SuppressWarnings("unchecked")
474+
List<BiConsumer<ClientEvent, AgentCard>> a2aEventConsumer =
475+
invocation.getArgument(1, List.class);
476+
477+
// Task creation
478+
Task initialTask =
479+
new Task.Builder()
480+
.id("t1")
481+
.contextId("c1")
482+
.status(new TaskStatus(TaskState.WORKING))
483+
.build();
484+
a2aEventConsumer.forEach(c -> c.accept(new TaskEvent(initialTask), agentCard));
485+
486+
// Stream output a piece of text (Artifact Update)
487+
TaskArtifactUpdateEvent chunkEvent =
488+
new TaskArtifactUpdateEvent.Builder()
489+
.taskId("t1")
490+
.contextId("c1")
491+
.artifact(
492+
new Artifact.Builder()
493+
.artifactId("a1")
494+
.name("mockArtifact")
495+
.parts(new TextPart("Hello A2A"))
496+
.build())
497+
.build();
498+
Task workingTask =
499+
new Task.Builder()
500+
.id("t1")
501+
.contextId("c1")
502+
.status(new TaskStatus(TaskState.WORKING))
503+
.artifacts(List.of(chunkEvent.getArtifact()))
504+
.build();
505+
a2aEventConsumer.forEach(
506+
c -> c.accept(new TaskUpdateEvent(workingTask, chunkEvent), agentCard));
507+
508+
// Task complete (Status Update - COMPLETED)
509+
Task completedTask =
510+
new Task.Builder()
511+
.id("t1")
512+
.contextId("c1")
513+
.status(new TaskStatus(TaskState.COMPLETED))
514+
.artifacts(List.of(chunkEvent.getArtifact()))
515+
.build();
516+
TaskStatusUpdateEvent completeEvent =
517+
new TaskStatusUpdateEvent(
518+
"t1",
519+
new TaskStatus(TaskState.COMPLETED),
520+
"c1",
521+
true,
522+
Map.of());
523+
a2aEventConsumer.forEach(
524+
c ->
525+
c.accept(
526+
new TaskUpdateEvent(completedTask, completeEvent),
527+
agentCard));
528+
529+
return null;
530+
};
531+
532+
doAnswer(mockTaskResponse)
533+
.when(a2aClient)
534+
.sendMessage(any(Message.class), anyList(), any());
535+
536+
agent.stream(Msg.builder().textContent("测试触发").build()).collectList().block();
537+
538+
assertEquals(1, preCount.get());
539+
assertEquals(1, chunkCount.get());
540+
assertEquals(1, postCount.get());
541+
}
542+
431543
private Answer<Void> mockSuccessMessage() {
432544
return invocationOnMock -> {
433545
@SuppressWarnings("unchecked")

0 commit comments

Comments
 (0)