Skip to content

Commit 7fbdc37

Browse files
ark-hand[bot]ark-hand[bot]
committed
fix(selfHosted): filter recovered tool results by session blockers
Sync-Source-Commit: 3d1df8fbeb78a4093bf8a337f8dda94fc1919fd9 Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.7.0
1 parent bd84806 commit 7fbdc37

6 files changed

Lines changed: 296 additions & 3 deletions

File tree

‎src/main/java/com/volcengine/ark/runtime/selfhosted/Event.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,25 @@ public String stopReasonType() {
138138
return stopReason instanceof String ? (String) stopReason : "";
139139
}
140140

141+
@SuppressWarnings("unchecked")
142+
public List<String> stopReasonEventIds() {
143+
if (!(stopReason instanceof Map)) {
144+
return new ArrayList<>();
145+
}
146+
Object rawIds = ((Map<String, Object>) stopReason).get("event_ids");
147+
if (!(rawIds instanceof List)) {
148+
return new ArrayList<>();
149+
}
150+
List<String> ids = new ArrayList<>();
151+
for (Object rawId : (List<Object>) rawIds) {
152+
String id = stringValue(rawId);
153+
if (!id.isEmpty()) {
154+
ids.add(id);
155+
}
156+
}
157+
return ids;
158+
}
159+
141160
public String callId() {
142161
if (toolUseId != null && !toolUseId.isEmpty()) {
143162
return toolUseId;

‎src/main/java/com/volcengine/ark/runtime/selfhosted/FileToolResultStore.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,14 @@ public void markSent(String callId) throws IOException {
122122
write(record);
123123
}
124124

125+
public void discard(String callId) throws IOException {
126+
if (callId == null || callId.isEmpty()) {
127+
throw new IllegalArgumentException("call id must not be empty");
128+
}
129+
Files.deleteIfExists(path(callId));
130+
syncDirectory();
131+
}
132+
125133
private Map<String, Object> read(String callId) throws IOException {
126134
Path path = path(callId);
127135
if (!Files.exists(path)) {

‎src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedConstants.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ public final class SelfHostedConstants {
1313
public static final String EVENT_TYPE_USER_TOOL_RESULT = "user.tool_result";
1414
public static final String EVENT_TYPE_USER_CUSTOM_TOOL_RESULT = "user.custom_tool_result";
1515
public static final String EVENT_TYPE_SESSION_STATUS_IDLE = "session.status_idle";
16+
public static final String EVENT_TYPE_SESSION_STATUS_RUNNING = "session.status_running";
17+
public static final String EVENT_TYPE_SESSION_STATUS_RESCHEDULED = "session.status_rescheduled";
1618
public static final String EVENT_TYPE_SESSION_STATUS_TERMINATED = "session.status_terminated";
1719
public static final String EVENT_TYPE_SESSION_DELETED = "session.deleted";
1820

@@ -25,6 +27,7 @@ public final class SelfHostedConstants {
2527

2628
public static final String EVENT_LIST_ORDER_ASC = "asc";
2729
public static final String SESSION_STOP_REASON_END_TURN = "end_turn";
30+
public static final String SESSION_STOP_REASON_REQUIRES_ACTION = "requires_action";
2831

2932
public static final long DEFAULT_MAX_IDLE_MILLIS = 60000L;
3033
public static final long DEFAULT_TOOL_TIMEOUT_MILLIS = 120000L;

‎src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java‎

Lines changed: 85 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,9 @@ public List<ToolCallResult> run() throws IOException {
6666
if (options.resultStore != null) {
6767
FileToolResultStore.RecoverResult recovered = options.resultStore.recover();
6868
state.pendingResults.putAll(recovered.getPending());
69+
for (String callId : recovered.getPending().keySet()) {
70+
state.recoveredResults.put(callId, Boolean.TRUE);
71+
}
6972
state.processed.putAll(recovered.getProcessed());
7073
state.answered.putAll(recovered.getProcessed());
7174
}
@@ -202,8 +205,8 @@ private void reconcileOnce(boolean reconcile) throws IOException {
202205

203206
private void consumeList() throws IOException {
204207
while (!isClosed()) {
205-
flushResults();
206208
reconcile(false);
209+
flushResults();
207210
if (idleExpired()) {
208211
throw new IdleTimeoutException();
209212
}
@@ -231,6 +234,7 @@ private void processListedEvents(List<Event> events, boolean reconcile) throws I
231234
if (!reconcile && !seenNow) {
232235
continue;
233236
}
237+
observeSessionState(event);
234238
if (seenNow && !SelfHostedConstants.EVENT_TYPE_USER_TOOL_CONFIRMATION.equals(event.getType())) {
235239
touchedIdle = true;
236240
lastWasEndTurn = SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_IDLE.equals(event.getType())
@@ -254,11 +258,12 @@ private void processListedEvents(List<Event> events, boolean reconcile) throws I
254258
throw new SessionTerminatedException();
255259
}
256260
}
261+
reconcileRecoveredResults();
257262
if (touchedIdle) {
258263
disarmIdle();
259264
}
260265
for (Event event : pending) {
261-
if (!isAnswered(event.callId())) {
266+
if (!isAnswered(event.callId()) && shouldHandleToolUse(event.callId())) {
262267
handleToolUse(event, SelfHostedConstants.EVENT_TYPE_AGENT_CUSTOM_TOOL_USE.equals(event.getType()));
263268
}
264269
}
@@ -288,6 +293,8 @@ private void handleStreamEvent(Event event) throws IOException {
288293
if (!markEventSeen(event)) {
289294
return;
290295
}
296+
observeSessionState(event);
297+
reconcileRecoveredResults();
291298
noteIdleEvent(event);
292299
handleEvent(event);
293300
}
@@ -316,6 +323,9 @@ private void handleToolUse(Event event, boolean custom) throws IOException {
316323
}
317324
Event pending = state.pendingResults.get(callId);
318325
if (pending != null) {
326+
if (state.recoveredResults.containsKey(callId)) {
327+
return;
328+
}
319329
sendResult(callId, event, custom, "", pending);
320330
return;
321331
}
@@ -468,6 +478,9 @@ private boolean retrySendEvent(Event event) {
468478

469479
private void flushResults() throws IOException {
470480
for (Map.Entry<String, Event> entry : new ArrayList<>(state.pendingResults.entrySet())) {
481+
if (state.recoveredResults.containsKey(entry.getKey())) {
482+
continue;
483+
}
471484
if (retrySendEvent(entry.getValue())) {
472485
markAnswered(entry.getKey());
473486
if (options.resultStore != null) {
@@ -486,6 +499,69 @@ private void flushResults() throws IOException {
486499
maybeArmPendingIdle();
487500
}
488501

502+
private void observeSessionState(Event event) {
503+
String type = event.getType();
504+
if (SelfHostedConstants.EVENT_TYPE_AGENT_TOOL_USE.equals(type)
505+
|| SelfHostedConstants.EVENT_TYPE_AGENT_CUSTOM_TOOL_USE.equals(type)) {
506+
String callId = event.callId();
507+
if (!callId.isEmpty()) {
508+
state.sessionToolUses.put(callId, Boolean.TRUE);
509+
state.toolUsesSinceStatus.put(callId, Boolean.TRUE);
510+
}
511+
return;
512+
}
513+
if (SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_IDLE.equals(type)) {
514+
state.blockingEventsKnown = true;
515+
state.blockingEventIds.clear();
516+
if (SelfHostedConstants.SESSION_STOP_REASON_REQUIRES_ACTION.equals(event.stopReasonType())) {
517+
for (String eventId : event.stopReasonEventIds()) {
518+
state.blockingEventIds.put(eventId, Boolean.TRUE);
519+
}
520+
}
521+
state.toolUsesSinceStatus.clear();
522+
return;
523+
}
524+
if (SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_RUNNING.equals(type)
525+
|| SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_RESCHEDULED.equals(type)) {
526+
state.blockingEventsKnown = true;
527+
state.blockingEventIds.clear();
528+
state.toolUsesSinceStatus.clear();
529+
}
530+
}
531+
532+
private boolean shouldHandleToolUse(String callId) {
533+
if (!state.blockingEventsKnown) {
534+
return true;
535+
}
536+
return state.blockingEventIds.containsKey(callId) || state.toolUsesSinceStatus.containsKey(callId);
537+
}
538+
539+
private void reconcileRecoveredResults() {
540+
if (!state.blockingEventsKnown) {
541+
return;
542+
}
543+
for (String callId : new ArrayList<>(state.recoveredResults.keySet())) {
544+
if (state.blockingEventIds.containsKey(callId) && state.sessionToolUses.containsKey(callId)) {
545+
state.recoveredResults.remove(callId);
546+
continue;
547+
}
548+
if (state.toolUsesSinceStatus.containsKey(callId)) {
549+
continue;
550+
}
551+
state.recoveredResults.remove(callId);
552+
state.pendingResults.remove(callId);
553+
LOGGER.warning("discard stale recovered tool result tool_use_id=" + callId);
554+
if (options.resultStore != null) {
555+
try {
556+
options.resultStore.discard(callId);
557+
} catch (IOException error) {
558+
LOGGER.log(Level.WARNING, "discard persisted tool result failed tool_use_id=" + callId, error);
559+
}
560+
}
561+
}
562+
maybeArmPendingIdle();
563+
}
564+
489565
private boolean ownsTool(Event event, boolean custom) {
490566
return custom ? options.customTools.containsKey(event.getName()) : options.tools.has(event.getName());
491567
}
@@ -537,6 +613,7 @@ private void markAnswered(String callId) {
537613
state.answered.put(callId, Boolean.TRUE);
538614
state.processed.put(callId, Boolean.TRUE);
539615
state.pendingResults.remove(callId);
616+
state.recoveredResults.remove(callId);
540617
state.pendingAsk.remove(callId);
541618
state.externalTools.remove(callId);
542619
maybeArmPendingIdle();
@@ -565,7 +642,7 @@ private void releaseConfirmedToolUses() throws IOException {
565642
private boolean hasUnblockedOutstandingTool(List<Event> pending) {
566643
for (Event event : pending) {
567644
String callId = event.callId();
568-
if (callId.isEmpty() || isAnswered(callId)) {
645+
if (callId.isEmpty() || isAnswered(callId) || !shouldHandleToolUse(callId)) {
569646
continue;
570647
}
571648
if (state.pendingAsk.containsKey(callId) || state.pendingResults.containsKey(callId)) {
@@ -661,9 +738,14 @@ private static class State {
661738
Map<String, Boolean> seen = new LinkedHashMap<>();
662739
Map<String, Boolean> answered = new LinkedHashMap<>();
663740
Map<String, Event> pendingResults = new LinkedHashMap<>();
741+
Map<String, Boolean> recoveredResults = new LinkedHashMap<>();
664742
Map<String, Event> pendingAsk = new LinkedHashMap<>();
665743
Map<String, Event> confirmations = new LinkedHashMap<>();
666744
Map<String, Event> externalTools = new LinkedHashMap<>();
745+
Map<String, Boolean> sessionToolUses = new LinkedHashMap<>();
746+
Map<String, Boolean> toolUsesSinceStatus = new LinkedHashMap<>();
747+
Map<String, Boolean> blockingEventIds = new LinkedHashMap<>();
748+
boolean blockingEventsKnown;
667749
long idleArmedAt;
668750
boolean idleArmPending;
669751
}

‎src/test/java/com/volcengine/ark/runtime/selfhosted/FileToolResultStoreTest.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,4 +73,20 @@ public void sessionStoreSanitizesSessionId() throws Exception {
7373
}
7474
assertFalse(Files.exists(workdir.resolve("outside")));
7575
}
76+
77+
@Test
78+
public void discardRemovesRecoveredRecord() throws Exception {
79+
Path workdir = Files.createTempDirectory("ark-java-store-");
80+
FileToolResultStore store = new FileToolResultStore(workdir.toString(), "session-a");
81+
Map<String, Object> raw = new LinkedHashMap<>();
82+
raw.put("id", "call-1");
83+
raw.put("type", "agent.tool_use");
84+
raw.put("name", "bash");
85+
store.begin("call-1", Event.fromMap(raw));
86+
87+
store.discard("call-1");
88+
89+
assertTrue(store.recover().getPending().isEmpty());
90+
assertTrue(store.recover().getProcessed().isEmpty());
91+
}
7692
}

0 commit comments

Comments
 (0)