Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions src/main/java/com/volcengine/ark/runtime/selfhosted/Event.java
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,25 @@ public String stopReasonType() {
return stopReason instanceof String ? (String) stopReason : "";
}

@SuppressWarnings("unchecked")
public List<String> stopReasonEventIds() {
if (!(stopReason instanceof Map)) {
return new ArrayList<>();
}
Object rawIds = ((Map<String, Object>) stopReason).get("event_ids");
if (!(rawIds instanceof List)) {
return new ArrayList<>();
}
List<String> ids = new ArrayList<>();
for (Object rawId : (List<Object>) rawIds) {
String id = stringValue(rawId);
if (!id.isEmpty()) {
ids.add(id);
}
}
return ids;
}

public String callId() {
if (toolUseId != null && !toolUseId.isEmpty()) {
return toolUseId;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,14 @@ public void markSent(String callId) throws IOException {
write(record);
}

public void discard(String callId) throws IOException {
if (callId == null || callId.isEmpty()) {
throw new IllegalArgumentException("call id must not be empty");
}
Files.deleteIfExists(path(callId));
syncDirectory();
}

private Map<String, Object> read(String callId) throws IOException {
Path path = path(callId);
if (!Files.exists(path)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ public final class SelfHostedConstants {
public static final String EVENT_TYPE_USER_TOOL_RESULT = "user.tool_result";
public static final String EVENT_TYPE_USER_CUSTOM_TOOL_RESULT = "user.custom_tool_result";
public static final String EVENT_TYPE_SESSION_STATUS_IDLE = "session.status_idle";
public static final String EVENT_TYPE_SESSION_STATUS_RUNNING = "session.status_running";
public static final String EVENT_TYPE_SESSION_STATUS_RESCHEDULED = "session.status_rescheduled";
public static final String EVENT_TYPE_SESSION_STATUS_TERMINATED = "session.status_terminated";
public static final String EVENT_TYPE_SESSION_DELETED = "session.deleted";

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

public static final String EVENT_LIST_ORDER_ASC = "asc";
public static final String SESSION_STOP_REASON_END_TURN = "end_turn";
public static final String SESSION_STOP_REASON_REQUIRES_ACTION = "requires_action";

public static final long DEFAULT_MAX_IDLE_MILLIS = 60000L;
public static final long DEFAULT_TOOL_TIMEOUT_MILLIS = 120000L;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,9 @@ public List<ToolCallResult> run() throws IOException {
if (options.resultStore != null) {
FileToolResultStore.RecoverResult recovered = options.resultStore.recover();
state.pendingResults.putAll(recovered.getPending());
for (String callId : recovered.getPending().keySet()) {
state.recoveredResults.put(callId, Boolean.TRUE);
}
state.processed.putAll(recovered.getProcessed());
state.answered.putAll(recovered.getProcessed());
}
Expand Down Expand Up @@ -202,8 +205,8 @@ private void reconcileOnce(boolean reconcile) throws IOException {

private void consumeList() throws IOException {
while (!isClosed()) {
flushResults();
reconcile(false);
flushResults();
if (idleExpired()) {
throw new IdleTimeoutException();
}
Expand Down Expand Up @@ -231,6 +234,7 @@ private void processListedEvents(List<Event> events, boolean reconcile) throws I
if (!reconcile && !seenNow) {
continue;
}
observeSessionState(event);
if (seenNow && !SelfHostedConstants.EVENT_TYPE_USER_TOOL_CONFIRMATION.equals(event.getType())) {
touchedIdle = true;
lastWasEndTurn = SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_IDLE.equals(event.getType())
Expand All @@ -254,11 +258,12 @@ private void processListedEvents(List<Event> events, boolean reconcile) throws I
throw new SessionTerminatedException();
}
}
reconcileRecoveredResults();
if (touchedIdle) {
disarmIdle();
}
for (Event event : pending) {
if (!isAnswered(event.callId())) {
if (!isAnswered(event.callId()) && shouldHandleToolUse(event.callId())) {
handleToolUse(event, SelfHostedConstants.EVENT_TYPE_AGENT_CUSTOM_TOOL_USE.equals(event.getType()));
}
}
Expand Down Expand Up @@ -288,6 +293,8 @@ private void handleStreamEvent(Event event) throws IOException {
if (!markEventSeen(event)) {
return;
}
observeSessionState(event);
reconcileRecoveredResults();
noteIdleEvent(event);
handleEvent(event);
}
Expand Down Expand Up @@ -316,6 +323,9 @@ private void handleToolUse(Event event, boolean custom) throws IOException {
}
Event pending = state.pendingResults.get(callId);
if (pending != null) {
if (state.recoveredResults.containsKey(callId)) {
return;
}
sendResult(callId, event, custom, "", pending);
return;
}
Expand Down Expand Up @@ -468,6 +478,9 @@ private boolean retrySendEvent(Event event) {

private void flushResults() throws IOException {
for (Map.Entry<String, Event> entry : new ArrayList<>(state.pendingResults.entrySet())) {
if (state.recoveredResults.containsKey(entry.getKey())) {
continue;
}
if (retrySendEvent(entry.getValue())) {
markAnswered(entry.getKey());
if (options.resultStore != null) {
Expand All @@ -486,6 +499,69 @@ private void flushResults() throws IOException {
maybeArmPendingIdle();
}

private void observeSessionState(Event event) {
String type = event.getType();
if (SelfHostedConstants.EVENT_TYPE_AGENT_TOOL_USE.equals(type)
|| SelfHostedConstants.EVENT_TYPE_AGENT_CUSTOM_TOOL_USE.equals(type)) {
String callId = event.callId();
if (!callId.isEmpty()) {
state.sessionToolUses.put(callId, Boolean.TRUE);
state.toolUsesSinceStatus.put(callId, Boolean.TRUE);
}
return;
}
if (SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_IDLE.equals(type)) {
state.blockingEventsKnown = true;
state.blockingEventIds.clear();
if (SelfHostedConstants.SESSION_STOP_REASON_REQUIRES_ACTION.equals(event.stopReasonType())) {
for (String eventId : event.stopReasonEventIds()) {
state.blockingEventIds.put(eventId, Boolean.TRUE);
}
}
state.toolUsesSinceStatus.clear();
return;
}
if (SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_RUNNING.equals(type)
|| SelfHostedConstants.EVENT_TYPE_SESSION_STATUS_RESCHEDULED.equals(type)) {
state.blockingEventsKnown = true;
state.blockingEventIds.clear();
state.toolUsesSinceStatus.clear();
}
}

private boolean shouldHandleToolUse(String callId) {
if (!state.blockingEventsKnown) {
return true;
}
return state.blockingEventIds.containsKey(callId) || state.toolUsesSinceStatus.containsKey(callId);
}

private void reconcileRecoveredResults() {
if (!state.blockingEventsKnown) {
return;
}
for (String callId : new ArrayList<>(state.recoveredResults.keySet())) {
if (state.blockingEventIds.containsKey(callId) && state.sessionToolUses.containsKey(callId)) {
state.recoveredResults.remove(callId);
continue;
}
if (state.toolUsesSinceStatus.containsKey(callId)) {
continue;
}
state.recoveredResults.remove(callId);
state.pendingResults.remove(callId);
LOGGER.warning("discard stale recovered tool result tool_use_id=" + callId);
if (options.resultStore != null) {
try {
options.resultStore.discard(callId);
} catch (IOException error) {
LOGGER.log(Level.WARNING, "discard persisted tool result failed tool_use_id=" + callId, error);
}
}
}
maybeArmPendingIdle();
}

private boolean ownsTool(Event event, boolean custom) {
return custom ? options.customTools.containsKey(event.getName()) : options.tools.has(event.getName());
}
Expand Down Expand Up @@ -537,6 +613,7 @@ private void markAnswered(String callId) {
state.answered.put(callId, Boolean.TRUE);
state.processed.put(callId, Boolean.TRUE);
state.pendingResults.remove(callId);
state.recoveredResults.remove(callId);
state.pendingAsk.remove(callId);
state.externalTools.remove(callId);
maybeArmPendingIdle();
Expand Down Expand Up @@ -565,7 +642,7 @@ private void releaseConfirmedToolUses() throws IOException {
private boolean hasUnblockedOutstandingTool(List<Event> pending) {
for (Event event : pending) {
String callId = event.callId();
if (callId.isEmpty() || isAnswered(callId)) {
if (callId.isEmpty() || isAnswered(callId) || !shouldHandleToolUse(callId)) {
continue;
}
if (state.pendingAsk.containsKey(callId) || state.pendingResults.containsKey(callId)) {
Expand Down Expand Up @@ -661,9 +738,14 @@ private static class State {
Map<String, Boolean> seen = new LinkedHashMap<>();
Map<String, Boolean> answered = new LinkedHashMap<>();
Map<String, Event> pendingResults = new LinkedHashMap<>();
Map<String, Boolean> recoveredResults = new LinkedHashMap<>();
Map<String, Event> pendingAsk = new LinkedHashMap<>();
Map<String, Event> confirmations = new LinkedHashMap<>();
Map<String, Event> externalTools = new LinkedHashMap<>();
Map<String, Boolean> sessionToolUses = new LinkedHashMap<>();
Map<String, Boolean> toolUsesSinceStatus = new LinkedHashMap<>();
Map<String, Boolean> blockingEventIds = new LinkedHashMap<>();
boolean blockingEventsKnown;
long idleArmedAt;
boolean idleArmPending;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,4 +73,20 @@ public void sessionStoreSanitizesSessionId() throws Exception {
}
assertFalse(Files.exists(workdir.resolve("outside")));
}

@Test
public void discardRemovesRecoveredRecord() throws Exception {
Path workdir = Files.createTempDirectory("ark-java-store-");
FileToolResultStore store = new FileToolResultStore(workdir.toString(), "session-a");
Map<String, Object> raw = new LinkedHashMap<>();
raw.put("id", "call-1");
raw.put("type", "agent.tool_use");
raw.put("name", "bash");
store.begin("call-1", Event.fromMap(raw));

store.discard("call-1");

assertTrue(store.recover().getPending().isEmpty());
assertTrue(store.recover().getProcessed().isEmpty());
}
}
Loading
Loading