Skip to content
Closed
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
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ public class HarnessAgent implements Agent, AutoCloseable {
private final SkillCurator skillCurator;
private final SkillAuditLog skillAuditLog;
private final MemoryConfig memoryConfig;
private final MemoryMaintenanceMiddleware memoryMaintenanceMw;

/** The subagent middleware (either SubagentsMiddleware or DynamicSubagentsMiddleware). */
private final Object subagentMiddleware;
Expand Down Expand Up @@ -199,6 +200,7 @@ private HarnessAgent(
SkillCurator skillCurator,
SkillAuditLog skillAuditLog,
MemoryConfig memoryConfig,
MemoryMaintenanceMiddleware memoryMaintenanceMw,
Object subagentMiddleware,
DistributedStore distributedStore,
WorkspacePathNormalizer pathNormalizer) {
Expand All @@ -217,6 +219,7 @@ private HarnessAgent(
this.skillCurator = skillCurator;
this.skillAuditLog = skillAuditLog;
this.memoryConfig = memoryConfig != null ? memoryConfig : MemoryConfig.defaults();
this.memoryMaintenanceMw = memoryMaintenanceMw;
this.subagentMiddleware = subagentMiddleware;
this.distributedStore = distributedStore;
this.pathNormalizer = pathNormalizer;
Expand Down Expand Up @@ -375,14 +378,20 @@ public io.agentscope.core.permission.PermissionMode getPermissionMode(
@Override
public void close() {
try {
shutdownTaskRepository();
if (memoryMaintenanceMw != null) {
memoryMaintenanceMw.close();
}
} finally {
try {
if (ownedWorkspaceIndex != null) {
ownedWorkspaceIndex.close();
}
shutdownTaskRepository();
} finally {
delegate.close();
try {
if (ownedWorkspaceIndex != null) {
ownedWorkspaceIndex.close();
}
} finally {
delegate.close();
}
}
}
}
Expand Down Expand Up @@ -2130,6 +2139,7 @@ public HarnessAgent build() {
inner.middleware(new AtPathExpansionMiddleware(wsManager));
}
Model memoryModel = memoryConfig.model() != null ? memoryConfig.model() : model;
MemoryMaintenanceMiddleware memoryMaintenanceMw = null;
if (memoryModel != null && !disableMemoryHooks) {
IsolationScope effectiveIsolationScope = fsIsolationScope;

Expand All @@ -2155,14 +2165,15 @@ public HarnessAgent build() {
memoryModel,
effectiveConsolidationPrompt,
memoryConfig.consolidationMaxTokens());
inner.middleware(
memoryMaintenanceMw =
new MemoryMaintenanceMiddleware(
wsManager,
consolidator,
memoryConfig.dailyFileRetentionDays(),
memoryConfig.sessionRetentionDays(),
memoryConfig.consolidationMinGap(),
effectiveIsolationScope));
effectiveIsolationScope);
inner.middleware(memoryMaintenanceMw);
}
CompactionMiddleware compactionHook = null;
if (!disableCompaction && compactionConfig != null) {
Expand Down Expand Up @@ -2491,6 +2502,7 @@ public HarnessAgent build() {
pendingSkillCurator,
pendingSkillAuditLog,
memoryConfig,
memoryMaintenanceMw,
capturedSubagentMw,
distributedStore,
pathNormalizer);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,13 @@
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.model.ExecutionConfig;
import io.agentscope.core.model.Model;
import io.agentscope.harness.agent.filesystem.AbstractFilesystem;
import io.agentscope.harness.agent.filesystem.model.FileInfo;
import io.agentscope.harness.agent.filesystem.model.GlobResult;
import io.agentscope.harness.agent.workspace.WorkspaceManager;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Comparator;
Expand Down Expand Up @@ -61,6 +63,15 @@ public class MemoryConsolidator {
/** Hidden state file inside {@code memory/} tracking the last consolidation Instant. */
public static final String STATE_FILE = ".consolidation_state";

/**
* Upper bound on the consolidation LLM call. Consolidation runs on a shared
* {@code boundedElastic} worker thread (see {@code MemoryMaintenanceMiddleware}); without a
* timeout, a hung model provider would tie up that thread indefinitely. Reuses
* {@link ExecutionConfig#MODEL_DEFAULTS}'s timeout so this stays consistent with the
* standard per-call model timeout used elsewhere in the codebase.
*/
static final Duration CONSOLIDATION_TIMEOUT = ExecutionConfig.MODEL_DEFAULTS.getTimeout();

/**
* Default prompt for the consolidation step. Exposed publicly so callers can extend
* (e.g. append project-specific guidelines) when constructing
Expand Down Expand Up @@ -167,6 +178,7 @@ public Mono<Void> consolidate(RuntimeContext rc) {
.build());

return model.stream(messages, null, null)
.timeout(CONSOLIDATION_TIMEOUT)
.reduce(
new StringBuilder(),
(sb, chatResponse) -> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,21 +29,26 @@
import java.time.Duration;
import java.time.Instant;
import java.time.LocalDate;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

/**
* Middleware that performs periodic memory maintenance after each agent call.
*
* <p>Fires on the agent invocation completion (via {@code onAgent concatWith}, after
* {@link MemoryFlushMiddleware}) and is throttled by a configurable minimum gap so it
* does not run on every single call.
* <p>Fires on the agent invocation completion (via {@code onAgent doOnComplete}, after
* {@link MemoryFlushMiddleware}) in a genuinely detached, fire-and-forget fashion: the
* maintenance {@code Mono} is subscribed independently of the returned {@code Flux} rather
* than being concatenated onto it, so callers that wait for the response to complete (e.g.
* {@code blockLast()}, {@code takeLast(1)}) are not delayed by maintenance work. It is also
* throttled by a configurable minimum gap so it does not run on every single call.
*
* <p>Maintenance steps executed in order:
* <ol>
Expand Down Expand Up @@ -77,6 +82,19 @@ public class MemoryMaintenanceMiddleware implements HarnessRuntimeMiddleware {
private final Duration minGap;
private final IsolationScope isolationScope;

/** Upper bound {@link #close()} waits for outstanding fire-and-forget runs to drain. */
static final Duration CLOSE_AWAIT_TIMEOUT = Duration.ofSeconds(5);

/**
* Tracks the {@link Disposable} of every fire-and-forget maintenance subscription that has
* been scheduled but not yet finished, so {@link #close()} can wait for/dispose them instead
* of leaving them racing against teardown of the resources they read/write (e.g. a workspace
* directory being deleted).
*/
private final Set<Disposable> pending = ConcurrentHashMap.newKeySet();

private volatile boolean closed = false;

/**
* Process-wide per-isolation-key maintenance timestamps. Static so that the throttle window
* survives across {@code HarnessAgent.Builder.build()} calls — each rebuild creates a new
Expand Down Expand Up @@ -139,17 +157,62 @@ public Flux<AgentEvent> onAgent(
AgentInput input,
Function<AgentInput, Flux<AgentEvent>> next) {
final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty();
return next.apply(input)
.concatWith(
Mono.<AgentEvent>fromRunnable(() -> maybeRunMaintenance(rc))
.subscribeOn(Schedulers.boundedElastic())
.onErrorResume(
e -> {
log.warn(
"Memory maintenance failed: {}",
e.getMessage());
return Mono.empty();
}));
return next.apply(input).doOnComplete(() -> scheduleMaintenance(rc));
}

/**
* Fires the fire-and-forget maintenance {@code Mono} on {@code boundedElastic}, tracking its
* {@link Disposable} in {@link #pending} until it terminates so {@link #close()} can wait for
* it. No-ops once {@link #close()} has been called, so a call that races with shutdown
* doesn't spawn new untracked work.
*/
private void scheduleMaintenance(RuntimeContext rc) {
if (closed) {
return;
}
Disposable[] holder = new Disposable[1];
Disposable d =
Mono.fromRunnable(() -> maybeRunMaintenance(rc))
.subscribeOn(Schedulers.boundedElastic())
.onErrorResume(
e -> {
log.warn("Memory maintenance failed: {}", e.getMessage());
return Mono.empty();
})
.doFinally(
sig -> {
if (holder[0] != null) {
pending.remove(holder[0]);
}
})
.subscribe();
holder[0] = d;
pending.add(d);
}

/**
* Waits (bounded by {@link #CLOSE_AWAIT_TIMEOUT}) for outstanding fire-and-forget maintenance
* runs to finish, then disposes anything still outstanding. Intended to be called from
* {@code HarnessAgent#close()} so short-lived callers (tests using JUnit {@code @TempDir},
* CLI runs, etc.) don't tear down the workspace while a detached maintenance write is still
* in flight — see {@link MemoryMaintenanceMiddleware class docs} for why maintenance is
* detached in the first place.
*/
public void close() {
closed = true;
long deadline = System.nanoTime() + CLOSE_AWAIT_TIMEOUT.toNanos();
while (!pending.isEmpty() && System.nanoTime() < deadline) {
try {
Thread.sleep(20);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
break;
}
}
for (Disposable d : pending) {
d.dispose();
}
pending.clear();
}

private void maybeRunMaintenance(RuntimeContext rc) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
/*
* Copyright 2024-2026 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.agentscope.harness.agent.memory;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import io.agentscope.core.agent.RuntimeContext;
import io.agentscope.core.model.ChatResponse;
import io.agentscope.core.model.Model;
import io.agentscope.harness.agent.filesystem.remote.RemoteFilesystem;
import io.agentscope.harness.agent.filesystem.remote.store.InMemoryStore;
import io.agentscope.harness.agent.workspace.WorkspaceManager;
import java.nio.file.Path;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.TimeoutException;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import reactor.core.publisher.Flux;
import reactor.test.StepVerifier;

/**
* Verifies that {@link MemoryConsolidator#consolidate} does not hang forever when the model
* provider never responds. Without a timeout, this would tie up a shared {@code boundedElastic}
* worker thread indefinitely (see {@link io.agentscope.harness.agent.middleware
* .MemoryMaintenanceMiddleware}, which runs consolidation on that scheduler).
*/
class MemoryConsolidatorTimeoutTest {

@Test
void consolidate_timesOutInsteadOfHangingForever(@TempDir Path tmp) throws Exception {
InMemoryStore store = new InMemoryStore();
List<String> ns = List.of("test-ns");
RemoteFilesystem fs = new RemoteFilesystem(store, ns);

try (WorkspaceManager wsm = new WorkspaceManager(tmp, fs)) {
// Seed a fresh daily entry so consolidate() doesn't short-circuit with Mono.empty()
// before ever reaching the model call.
wsm.writeUtf8WorkspaceRelative(
RuntimeContext.empty(), "memory/2025-06-15.md", "Some daily notes");

Model hungModel = mock(Model.class);
when(hungModel.stream(anyList(), any(), any())).thenReturn(Flux.<ChatResponse>never());

MemoryConsolidator consolidator = new MemoryConsolidator(wsm, hungModel);

StepVerifier.withVirtualTime(() -> consolidator.consolidate(RuntimeContext.empty()))
.thenAwait(MemoryConsolidator.CONSOLIDATION_TIMEOUT.plus(Duration.ofSeconds(1)))
.expectError(TimeoutException.class)
.verify(Duration.ofSeconds(10));
}
}
}
Loading
Loading