Skip to content

Commit 444ca06

Browse files
HandleCodingHandleCoding
committed
fix: add bypass stop and new cancellation
1 parent 4bc9146 commit 444ca06

28 files changed

Lines changed: 1710 additions & 45 deletions

‎application/src/main/java/io/github/huskyagent/application/ChatResult.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public enum ErrorCode {
2323
AUTH_ERROR,
2424
SESSION_ERROR,
2525
LLM_ERROR,
26+
CANCELLED,
2627
INTERNAL_ERROR
2728
}
2829

@@ -41,4 +42,9 @@ public static ChatResult failure(String errorMessage) {
4142
public static ChatResult failure(String errorMessage, ErrorCode errorCode) {
4243
return new ChatResult(null, false, errorMessage, errorCode, null, false, null);
4344
}
45+
46+
public static ChatResult cancelled(String sessionId, String message) {
47+
return new ChatResult(null, false, message != null ? message : "Run cancelled",
48+
ErrorCode.CANCELLED, sessionId, false, null);
49+
}
4450
}

‎application/src/main/java/io/github/huskyagent/application/ReActAgentApp.java‎

Lines changed: 63 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import io.github.huskyagent.application.agent.ClarifyContext;
55
import io.github.huskyagent.application.agent.TextEvent;
66
import io.github.huskyagent.application.runtime.AgentRuntimeExecutor;
7+
import io.github.huskyagent.application.runtime.RunCancelledException;
8+
import io.github.huskyagent.application.runtime.RunHandle;
79
import io.github.huskyagent.application.runtime.RuntimeCallbacks;
810
import io.github.huskyagent.application.runtime.RuntimeExecutionRequest;
911
import io.github.huskyagent.application.session.GraphCacheKey;
@@ -73,6 +75,7 @@ public class ReActAgentApp implements AgentRuntimeExecutor {
7375
private final MultimodalMessageBuilder multimodalMessageBuilder;
7476
private final DynamicPromptSnapshotCache dynamicPromptSnapshotCache;
7577
private final ToolCallbackFactory toolCallbackFactory;
78+
private final io.github.huskyagent.application.runtime.SessionRunCoordinator runCoordinator;
7679

7780
/**
7881
* Caches compiled graphs by runtime-policy fingerprint so scenes, principals,
@@ -103,6 +106,13 @@ public ChatResult execute(RuntimeScope scope, AgentInput input, RuntimeCallbacks
103106
@Override
104107
public ChatResult execute(RuntimeScope scope, AgentInput input, RuntimeCallbacks callbacks,
105108
RuntimeExecutionRequest.PersistenceMode persistenceMode) {
109+
return execute(scope, input, callbacks, persistenceMode, null);
110+
}
111+
112+
@Override
113+
public ChatResult execute(RuntimeScope scope, AgentInput input, RuntimeCallbacks callbacks,
114+
RuntimeExecutionRequest.PersistenceMode persistenceMode,
115+
RunHandle runHandle) {
106116
RuntimeCallbacks executionCallbacks = callbacks != null ? callbacks : RuntimeCallbacks.NOOP;
107117
RuntimeExecutionRequest.PersistenceMode mode = persistenceMode != null
108118
? persistenceMode
@@ -121,14 +131,21 @@ public ChatResult execute(RuntimeScope scope, AgentInput input, RuntimeCallbacks
121131
currentCheckpointId(graph, sid));
122132
}
123133
String turnId = UUID.randomUUID().toString();
124-
RunnableConfig config = buildConfig(sid, turnId, scope);
134+
RunnableConfig config = buildConfig(sid, turnId, scope, runHandle);
125135
try {
126136
return runWithInterruptLoop(scope, graph, config, buildInputs(input), sid,
127-
executionCallbacks, stateless);
137+
executionCallbacks, stateless, runHandle);
128138
} finally {
129139
dynamicPromptSnapshotCache.clearTurn(sid, turnId);
130140
}
141+
} catch (RunCancelledException e) {
142+
log.info("Graph execution cancelled: sessionId={}", sid);
143+
return ChatResult.cancelled(sid, e.getMessage());
131144
} catch (Exception e) {
145+
if (isCancellation(e, runHandle)) {
146+
log.info("Graph execution cancelled: sessionId={}", sid);
147+
return ChatResult.cancelled(sid, "Run cancelled");
148+
}
132149
log.error("Graph execution failed: sessionId={}", sid, e);
133150
return ChatResult.failure(e.getMessage());
134151
}
@@ -159,16 +176,19 @@ private ChatResult runWithInterruptLoop(
159176
Map<String, Object> inputs,
160177
String sessionId,
161178
RuntimeCallbacks callbacks,
162-
boolean stateless) throws Exception {
179+
boolean stateless,
180+
RunHandle runHandle) throws Exception {
163181

164182
Map<String, Object> currentInputs = inputs;
165183
ReActAgentState finalState = null;
166184
int resumeCount = 0;
167185
boolean hasModelOutput = false;
168186

169187
while (true) {
188+
throwIfCancelled(runHandle);
170189
var generator = graph.stream(currentInputs, config);
171190
for (NodeOutput<ReActAgentState> step : generator) {
191+
throwIfCancelled(runHandle);
172192
finalState = step.state();
173193
if (AgentGraph.NODE_MODEL.equals(step.node())) {
174194
hasModelOutput = true;
@@ -179,6 +199,7 @@ private ChatResult runWithInterruptLoop(
179199
log.debug("[loop] graphResult type={}", graphResult.type());
180200

181201
if (!graphResult.isInterruptionMetadata()) break;
202+
throwIfCancelled(runHandle);
182203

183204
if (resumeCount++ > 50) {
184205
log.warn("Interrupt resume loop exceeded the limit; forcing exit");
@@ -196,15 +217,41 @@ private ChatResult runWithInterruptLoop(
196217
callbacks.approval(scope, buildApprovalContext(
197218
graph, config, sessionId, metadata, finalState));
198219
}
220+
throwIfCancelled(runHandle);
199221
currentInputs = null;
200222
}
201223

224+
throwIfCancelled(runHandle);
202225
if (!stateless) {
203226
recordProviderTokenUsage(sessionId, finalState);
204227
compactActiveCheckpointIfNeeded(scope, graph, config, finalState);
205228
}
206229
return handleFinalState(sessionId, finalState, hasModelOutput, stateless);
207230
}
231+
private void throwIfCancelled(RunHandle runHandle) {
232+
if (runHandle != null && (Thread.currentThread().isInterrupted() || runCoordinator.isCancelled(runHandle))) {
233+
throw new RunCancelledException(runHandle.sessionId());
234+
}
235+
}
236+
237+
private boolean isCancellation(Throwable error, RunHandle runHandle) {
238+
return runHandle != null
239+
&& (runCoordinator.isCancelled(runHandle)
240+
|| Thread.currentThread().isInterrupted()
241+
|| hasCause(error, InterruptedException.class)
242+
|| hasCause(error, java.util.concurrent.CancellationException.class));
243+
}
244+
245+
private boolean hasCause(Throwable error, Class<? extends Throwable> type) {
246+
Throwable current = error;
247+
while (current != null) {
248+
if (type.isInstance(current)) {
249+
return true;
250+
}
251+
current = current.getCause();
252+
}
253+
return false;
254+
}
208255

209256

210257
private CompiledGraph<ReActAgentState> buildStatelessGraph(RuntimeScope scope) throws Exception {
@@ -278,17 +325,21 @@ private void initSession(String sessionId) {
278325
}
279326
}
280327

281-
private RunnableConfig buildConfig(String sessionId, String turnId, RuntimeScope scope) {
282-
return RunnableConfig.builder()
328+
private RunnableConfig buildConfig(String sessionId, String turnId, RuntimeScope scope, RunHandle runHandle) {
329+
var builder = RunnableConfig.builder()
283330
.threadId(sessionId)
284331
.putMetadata(DYNAMIC_PROMPT_TURN_ID_METADATA, turnId)
285-
.putMetadata(RequestToolContext.METADATA_KEY, buildRequestToolContext(sessionId, scope))
332+
.putMetadata(RequestToolContext.METADATA_KEY, buildRequestToolContext(sessionId, scope, runHandle != null ? runHandle.toDomain() : null))
286333
.putMetadata("channelIdentity", scope.getChannelIdentity())
287-
.putMetadata("principal", scope.getPrincipal())
288-
.build();
334+
.putMetadata("principal", scope.getPrincipal());
335+
if (runHandle != null) {
336+
builder.putMetadata(RunHandle.METADATA_KEY, runHandle);
337+
}
338+
return builder.build();
289339
}
290340

291-
private RequestToolContext buildRequestToolContext(String sessionId, RuntimeScope scope) {
341+
private RequestToolContext buildRequestToolContext(String sessionId, RuntimeScope scope,
342+
io.github.huskyagent.domain.runtime.RunHandle runHandle) {
292343
var runtimePolicy = scope.getRuntimePolicy();
293344
var capabilityView = runtimePolicy.getCapabilityView();
294345
var toolDefinitions = capabilityView.getVisibleTools();
@@ -300,7 +351,9 @@ private RequestToolContext buildRequestToolContext(String sessionId, RuntimeScop
300351
capabilityView.getVisibleSkillNames(),
301352
capabilityView.getVisiblePromptSections());
302353
return RequestToolContext.of(toolDefinitions,
303-
toolCallbackFactory.build(toolDefinitions, sessionId, executionContext));
354+
toolCallbackFactory.build(toolDefinitions, sessionId, executionContext),
355+
runHandle,
356+
runCoordinator);
304357
}
305358

306359
private Map<String, Object> buildInputs(AgentInput input) {

‎application/src/main/java/io/github/huskyagent/application/channel/BypassCommandPolicy.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package io.github.huskyagent.application.channel;
2+
3+
import org.springframework.stereotype.Component;
4+
5+
@Component
6+
public class BypassCommandPolicy {
7+
8+
public CommandExecutionMode modeFor(ChannelCommand command) {
9+
if (command == null) {
10+
return CommandExecutionMode.NORMAL_QUEUED;
11+
}
12+
return switch (command.name()) {
13+
case "stop" -> CommandExecutionMode.BYPASS_CANCEL_ACTIVE;
14+
case "new", "newsession", "new-session" -> CommandExecutionMode.BYPASS_REPLACE_ACTIVE_AND_PENDING;
15+
default -> CommandExecutionMode.NORMAL_QUEUED;
16+
};
17+
}
18+
}

‎application/src/main/java/io/github/huskyagent/application/channel/ChannelCommandService.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ public OutboundMessage execute(ChannelCommand command, InboundMessage inbound, S
1919
String name = command.name();
2020
return switch (name) {
2121
case "new", "newsession", "new-session" -> newSession(inbound, sceneId);
22+
case "stop" -> reply(inbound, null, "Stop is handled by the runtime bypass path.");
2223
case "session", "session-info" -> currentSession(inbound, sceneId);
2324
case "help" -> help(inbound);
2425
default -> unknown(command, inbound);
@@ -27,7 +28,7 @@ public OutboundMessage execute(ChannelCommand command, InboundMessage inbound, S
2728

2829
public boolean supports(ChannelCommand command) {
2930
return switch (command.name()) {
30-
case "new", "newsession", "new-session", "session", "session-info", "help" -> true;
31+
case "new", "newsession", "new-session", "stop", "session", "session-info", "help" -> true;
3132
default -> false;
3233
};
3334
}
@@ -54,7 +55,7 @@ private OutboundMessage currentSession(InboundMessage inbound, String sceneId) {
5455
}
5556

5657
private OutboundMessage help(InboundMessage inbound) {
57-
return reply(inbound, null, "Available commands:\n/new Create a new session\n/session Show current session\n/help Show help");
58+
return reply(inbound, null, "Available commands:\n/new Create a new session\n/stop Stop current run\n/session Show current session\n/help Show help");
5859
}
5960

6061
private OutboundMessage unknown(ChannelCommand command, InboundMessage inbound) {

‎application/src/main/java/io/github/huskyagent/application/channel/ChannelRuntimeService.java‎

Lines changed: 111 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,33 +6,141 @@
66
import io.github.huskyagent.application.runtime.RuntimeExecutionRequest;
77
import io.github.huskyagent.application.runtime.RuntimeExecutionResult;
88
import io.github.huskyagent.application.runtime.RuntimeExecutionService;
9+
import io.github.huskyagent.application.runtime.StopResult;
10+
import io.github.huskyagent.application.session.RuntimeScope;
11+
import io.github.huskyagent.application.session.SessionResolver;
912
import io.github.huskyagent.infra.channel.InboundMessage;
10-
import lombok.RequiredArgsConstructor;
13+
import io.github.huskyagent.infra.channel.OutboundMessage;
14+
import org.springframework.beans.factory.annotation.Autowired;
1115
import org.springframework.stereotype.Service;
1216

17+
import java.util.Optional;
1318
import java.util.concurrent.CompletableFuture;
1419
import java.util.concurrent.Executor;
1520

1621
@Service
17-
@RequiredArgsConstructor
1822
public class ChannelRuntimeService {
1923

2024
private final RuntimeExecutionService runtimeExecutionService;
2125
private final ChannelInboundQueue inboundQueue;
2226
private final ChannelRuntimeQueueKeyFactory queueKeyFactory;
2327
private final ChannelSceneRouter sceneRouter;
28+
private final ChannelCommandParser commandParser;
29+
private final BypassCommandPolicy bypassCommandPolicy;
30+
private final SessionResolver sessionResolver;
31+
32+
@Autowired
33+
public ChannelRuntimeService(RuntimeExecutionService runtimeExecutionService,
34+
ChannelInboundQueue inboundQueue,
35+
ChannelRuntimeQueueKeyFactory queueKeyFactory,
36+
ChannelSceneRouter sceneRouter,
37+
ChannelCommandParser commandParser,
38+
BypassCommandPolicy bypassCommandPolicy,
39+
SessionResolver sessionResolver) {
40+
this.runtimeExecutionService = runtimeExecutionService;
41+
this.inboundQueue = inboundQueue;
42+
this.queueKeyFactory = queueKeyFactory;
43+
this.sceneRouter = sceneRouter;
44+
this.commandParser = commandParser;
45+
this.bypassCommandPolicy = bypassCommandPolicy;
46+
this.sessionResolver = sessionResolver;
47+
}
48+
49+
public ChannelRuntimeService(RuntimeExecutionService runtimeExecutionService,
50+
ChannelInboundQueue inboundQueue,
51+
ChannelRuntimeQueueKeyFactory queueKeyFactory,
52+
ChannelSceneRouter sceneRouter) {
53+
this(runtimeExecutionService, inboundQueue, queueKeyFactory, sceneRouter,
54+
inbound -> Optional.empty(), new BypassCommandPolicy(), null);
55+
}
2456

2557
public CompletableFuture<ChatResult> handleInboundAsync(InboundMessage inbound, ChannelAdapter adapter, Executor executor) {
2658
EffectiveChannelRoute route = sceneRouter.resolve(inbound);
2759
String queueKey = queueKeyFactory.keyFor(inbound, route);
28-
return inboundQueue.enqueue(queueKey, () -> handleInbound(inbound, adapter, route), executor);
60+
Optional<ChannelCommand> command = parseCommand(inbound);
61+
CommandExecutionMode mode = command.map(bypassCommandPolicy::modeFor).orElse(CommandExecutionMode.NORMAL_QUEUED);
62+
if (mode != CommandExecutionMode.NORMAL_QUEUED) {
63+
try {
64+
return CompletableFuture.completedFuture(handleBypass(command.orElseThrow(), mode, inbound, adapter, route, queueKey));
65+
} catch (Exception e) {
66+
return CompletableFuture.completedFuture(ChatResult.failure(e.getMessage()));
67+
}
68+
}
69+
long generation = runtimeExecutionService.runCoordinator().currentQueueGeneration(queueKey);
70+
return inboundQueue.enqueue(queueKey, () -> {
71+
if (!runtimeExecutionService.runCoordinator().isQueueGenerationCurrent(queueKey, generation)) {
72+
return ChatResult.cancelled(inbound.getRequestedSessionId(), "Queued request superseded");
73+
}
74+
return handleInbound(inbound, adapter, route);
75+
}, executor);
2976
}
3077

3178
public ChatResult handleInbound(InboundMessage inbound, ChannelAdapter adapter) {
3279
EffectiveChannelRoute route = sceneRouter.resolve(inbound);
3380
return handleInbound(inbound, adapter, route);
3481
}
3582

83+
private Optional<ChannelCommand> parseCommand(InboundMessage inbound) {
84+
return inbound != null && inbound.getText() != null && !inbound.getText().isBlank()
85+
? commandParser.parse(inbound)
86+
: Optional.empty();
87+
}
88+
89+
private ChatResult handleBypass(ChannelCommand command, CommandExecutionMode mode, InboundMessage inbound,
90+
ChannelAdapter adapter, EffectiveChannelRoute route, String queueKey) {
91+
return switch (mode) {
92+
case BYPASS_CANCEL_ACTIVE -> stopActiveRun(inbound, adapter, route);
93+
case BYPASS_REPLACE_ACTIVE_AND_PENDING -> newSession(inbound, adapter, route, queueKey);
94+
case NORMAL_QUEUED -> handleInbound(inbound, adapter, route);
95+
};
96+
}
97+
98+
private ChatResult stopActiveRun(InboundMessage inbound, ChannelAdapter adapter, EffectiveChannelRoute route) {
99+
Optional<String> sessionId = activeSessionId(inbound, route);
100+
StopResult stopResult = sessionId
101+
.map(id -> runtimeExecutionService.interruptSession(id, "channel_stop"))
102+
.orElse(StopResult.none(null, "channel_stop"));
103+
String text = stopResult.hadActiveRun()
104+
? "Stopped current run."
105+
: "No active run to stop.";
106+
adapter.send(reply(inbound, sessionId.orElse(null), text));
107+
return ChatResult.success(text, sessionId.orElse(null), false);
108+
}
109+
110+
private ChatResult newSession(InboundMessage inbound, ChannelAdapter adapter, EffectiveChannelRoute route, String queueKey) {
111+
runtimeExecutionService.runCoordinator().bumpQueueGeneration(queueKey);
112+
activeSessionId(inbound, route)
113+
.ifPresent(id -> runtimeExecutionService.expireSessionRun(id, "channel_new_session"));
114+
if (sessionResolver == null) {
115+
ChatResult result = ChatResult.failure("Session resolver is not available");
116+
adapter.send(reply(inbound, null, result.errorMessage()));
117+
return result;
118+
}
119+
RuntimeScope scope = sessionResolver.createSession(inbound.getPrincipal(), inbound.getChannelIdentity(), route.sceneId());
120+
String text = "Created new session: " + scope.getSessionId();
121+
adapter.send(reply(inbound, scope.getSessionId(), text));
122+
return ChatResult.success(text, scope.getSessionId(), false);
123+
}
124+
125+
private Optional<String> activeSessionId(InboundMessage inbound, EffectiveChannelRoute route) {
126+
if (inbound.getRequestedSessionId() != null && !inbound.getRequestedSessionId().isBlank()) {
127+
return Optional.of(inbound.getRequestedSessionId());
128+
}
129+
return sessionResolver != null
130+
? sessionResolver.findActiveSessionId(inbound.getPrincipal(), inbound.getChannelIdentity(), route.sceneId())
131+
: Optional.empty();
132+
}
133+
134+
private OutboundMessage reply(InboundMessage inbound, String sessionId, String text) {
135+
return OutboundMessage.builder()
136+
.kind(OutboundMessage.Kind.TEXT)
137+
.sessionId(sessionId)
138+
.channelIdentity(inbound.getChannelIdentity())
139+
.replyTarget(inbound.getReplyTarget())
140+
.text(text)
141+
.build();
142+
}
143+
36144
private ChatResult handleInbound(InboundMessage inbound, ChannelAdapter adapter, EffectiveChannelRoute route) {
37145
RuntimeExecutionResult result = runtimeExecutionService.execute(RuntimeExecutionRequest.builder()
38146
.inbound(inbound)

‎application/src/main/java/io/github/huskyagent/application/channel/CommandExecutionMode.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package io.github.huskyagent.application.channel;
2+
3+
public enum CommandExecutionMode {
4+
NORMAL_QUEUED,
5+
BYPASS_CANCEL_ACTIVE,
6+
BYPASS_REPLACE_ACTIVE_AND_PENDING
7+
}

‎application/src/main/java/io/github/huskyagent/application/runtime/AgentRuntimeExecutor.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,10 @@ default ChatResult execute(RuntimeScope scope, AgentInput input, RuntimeCallback
1717
}
1818
return execute(scope, input, callbacks);
1919
}
20+
21+
default ChatResult execute(RuntimeScope scope, AgentInput input, RuntimeCallbacks callbacks,
22+
RuntimeExecutionRequest.PersistenceMode persistenceMode,
23+
RunHandle runHandle) {
24+
return execute(scope, input, callbacks, persistenceMode);
25+
}
2026
}

0 commit comments

Comments
 (0)