CompactionMiddleware.java

/*
 * 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.middleware;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectReader;
import io.agentscope.core.ReActAgent;
import io.agentscope.core.agent.Agent;
import io.agentscope.core.agent.RuntimeContext;
import io.agentscope.core.event.AgentEvent;
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
import io.agentscope.core.middleware.ModelCallInput;
import io.agentscope.core.middleware.ReasoningInput;
import io.agentscope.core.model.Model;
import io.agentscope.core.session.SessionLogException;
import io.agentscope.core.state.AgentState;
import io.agentscope.core.util.ExceptionUtils;
import io.agentscope.core.util.JsonException;
import io.agentscope.core.util.JsonUtils;
import io.agentscope.harness.agent.memory.MemoryFlushManager;
import io.agentscope.harness.agent.memory.compaction.CompactionConfig;
import io.agentscope.harness.agent.memory.compaction.ConversationCompactor;
import io.agentscope.harness.agent.memory.compaction.TokenCounterUtil;
import io.agentscope.harness.agent.workspace.WorkspaceManager;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.List;
import java.util.Locale;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.function.Function;
import java.util.function.ToLongFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Exceptions;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * Middleware that performs conversation compaction before each LLM reasoning call.
 *
 * <p>Fires on {@link #onReasoning}. When the compaction threshold is exceeded:
 * <ol>
 *   <li>Long-term memories are flushed from the prefix via {@link MemoryFlushManager}.</li>
 *   <li>The native session log retains the complete conversation before compaction.</li>
 *   <li>The prefix is distilled into a structured summary via one LLM call.</li>
 *   <li>The agent's working {@link AgentState#contextMutable() context} is replaced with
 *       {@code [summaryMsg] + preservedTail}.</li>
 *   <li>The downstream {@link ReasoningInput} is rebuilt with
 *       {@code [systemMsg] + [summaryMsg] + preservedTail}.</li>
 * </ol>
 *
 * <p>When {@link CompactionConfig#getTriggerTokens()} is 0 (dynamic mode, the default), the
 * effective trigger threshold is computed as {@code model.getContextWindowSize() - reserved}.
 * If the model does not report its context window, falls back to
 * {@link CompactionConfig#FALLBACK_TRIGGER_TOKENS}.
 */
public class CompactionMiddleware implements HarnessRuntimeMiddleware {

    private static final Logger log = LoggerFactory.getLogger(CompactionMiddleware.class);

    /**
     * Reads serialized histories as exact decimal trees: restored copies keep their values
     * across Integer/Double/BigDecimal coercion, while edits to decimal metadata beyond
     * double precision still compare unequal. Numbers compare by exact value only, so
     * scale-only and notation-only changes ("1.10" vs "1.1", "1e2" vs "100") count as
     * unchanged.
     */
    private static final ObjectReader HISTORY_VALUES =
            new ObjectMapper().reader().with(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS);

    private final WorkspaceManager workspaceManager;
    private final Model model;
    private final CompactionConfig config;

    public CompactionMiddleware(
            WorkspaceManager workspaceManager, Model model, CompactionConfig config) {
        this.workspaceManager = workspaceManager;
        this.model = model;
        this.config = config;
    }

    /** Narrow declaration: subclasses overriding more hooks must extend this set. */
    @Override
    public Set<ExtensionPoint> activePoints() {
        return EnumSet.of(ExtensionPoint.ON_REASONING);
    }

    @Override
    public Flux<AgentEvent> onReasoning(
            Agent agent,
            RuntimeContext ctx,
            ReasoningInput input,
            Function<ReasoningInput, Flux<AgentEvent>> next) {
        Model executionModel = agent instanceof ReActAgent react ? react.getModel() : model;
        if (executionModel == null) executionModel = model;
        int window = executionModel.getContextWindowSize();
        int budget =
                window > 0
                        ? Math.max(1, window - config.getReserved())
                        : CompactionConfig.FALLBACK_TRIGGER_TOKENS;
        return prepare(
                        agent,
                        ctx,
                        new ModelCallInput(
                                input.messages(), input.tools(), input.options(), executionModel),
                        budget)
                .flatMapMany(
                        prepared ->
                                next.apply(
                                        new ReasoningInput(
                                                prepared.messages(),
                                                prepared.tools(),
                                                prepared.options())));
    }

    /** Retry provider context overflow inside the same execution, before its recorder is closed. */
    public HarnessRuntimeMiddleware overflowRecovery() {
        return new HarnessRuntimeMiddleware() {
            @Override
            public Flux<AgentEvent> onReasoning(
                    Agent agent,
                    RuntimeContext ctx,
                    ReasoningInput input,
                    Function<ReasoningInput, Flux<AgentEvent>> next) {
                if (!(agent instanceof ReActAgent react) || !react.sessionLogEnabled())
                    return next.apply(input);
                AgentState initialState = RuntimeContext.resolveAgentState(ctx, agent);
                List<Msg> initialHistory =
                        initialState == null ? List.of() : List.copyOf(initialState.getContext());
                return next.apply(input)
                        .onErrorResume(
                                error -> {
                                    if (ExceptionUtils.containsInterruptedException(error)
                                            || SessionLogException.causedBy(error)
                                            || !isContextOverflowError(error))
                                        return Flux.error(error);
                                    AgentState state = RuntimeContext.resolveAgentState(ctx, agent);
                                    if (state == null || state.getContext().isEmpty())
                                        return Flux.error(error);
                                    List<Msg> before = List.copyOf(state.getContext());
                                    // The final request preparer may already have compacted state
                                    // before the provider rejected the request. Retry from that
                                    // latest history.
                                    var retryMessages = new ArrayList<Msg>();
                                    input.messages().stream()
                                            .filter(
                                                    message ->
                                                            message.getRole() == MsgRole.SYSTEM
                                                                    || (message.isRequestProjection(
                                                                                    initialHistory)
                                                                            && message
                                                                                    .isRequestProjection(
                                                                                            before)))
                                            .forEach(retryMessages::add);
                                    retryMessages.addAll(before);
                                    var forced =
                                            new CompactionMiddleware(
                                                    workspaceManager,
                                                    model,
                                                    CompactionConfig.builder()
                                                            .triggerMessages(1)
                                                            .keepMessages(1)
                                                            .keepTokens(0)
                                                            .flushBeforeCompact(
                                                                    config.isFlushBeforeCompact())
                                                            .build());
                                    return forced.prepareCandidate(
                                                    agent,
                                                    ctx,
                                                    new ModelCallInput(
                                                            retryMessages,
                                                            input.tools(),
                                                            input.options(),
                                                            react.getModel()),
                                                    CompactionConfig.FALLBACK_TRIGGER_TOKENS,
                                                    before)
                                            .flatMapMany(
                                                    candidate -> {
                                                        if (candidate.history().equals(before)
                                                                || !state.getContext()
                                                                        .equals(before))
                                                            return Flux.error(error);
                                                        applyToContext(state, candidate.history());
                                                        // next creates a new model call; the turn,
                                                        // execution and writer stay unchanged.
                                                        return next.apply(
                                                                        new ReasoningInput(
                                                                                candidate
                                                                                        .input()
                                                                                        .messages(),
                                                                                candidate
                                                                                        .input()
                                                                                        .tools(),
                                                                                candidate
                                                                                        .input()
                                                                                        .options()))
                                                                .onErrorResume(
                                                                        retryError ->
                                                                                Flux.error(
                                                                                        associateFirstOverflow(
                                                                                                error,
                                                                                                retryError)));
                                                    });
                                });
            }
        };
    }

    /**
     * Keeps the first provider overflow associated with a failed single retry without replacing the
     * retry's own error: as the primary cause when the cause slot is free, otherwise as a
     * suppressed exception. Cancellations, fatal errors and interruptions are never rewritten so
     * Reactor cancellation semantics survive.
     */
    private static Throwable associateFirstOverflow(Throwable first, Throwable retryError) {
        if (retryError == first
                || retryError.getCause() == first
                || retryError instanceof CancellationException
                || Exceptions.isCancel(retryError)
                || Exceptions.isFatal(retryError)
                || ExceptionUtils.containsInterruptedException(retryError)) return retryError;
        for (Throwable suppressed : retryError.getSuppressed()) {
            if (suppressed == first) return retryError;
        }
        if (retryError.getCause() == null) {
            try {
                retryError.initCause(first);
                return retryError;
            } catch (IllegalStateException alreadyInitialized) {
                // Providers build errors with an explicit null cause; that slot rejects
                // initCause(), so fall through and attach the overflow as suppressed.
            }
        }
        retryError.addSuppressed(first);
        return retryError;
    }

    public static boolean isContextOverflowError(Throwable error) {
        String message = error.getMessage();
        if (message == null) return false;
        String lower = message.toLowerCase(Locale.ROOT);
        return lower.contains("context_length_exceeded")
                || lower.contains("context length")
                || lower.contains("maximum context")
                || lower.contains("token limit")
                || lower.contains("too many tokens")
                || lower.contains("exceeds the model's maximum")
                || lower.contains("reduce the length");
    }

    /** Compact at the final model boundary, using the actual request's conversation allowance. */
    public Mono<ModelCallInput> prepare(
            Agent agent, RuntimeContext ctx, ModelCallInput input, int conversationBudget) {
        return Mono.defer(
                () -> {
                    AgentState state = RuntimeContext.resolveAgentState(ctx, agent);
                    List<Msg> original =
                            state == null ? List.of() : List.copyOf(state.contextMutable());
                    return prepareCandidate(agent, ctx, input, conversationBudget, original)
                            .map(
                                    candidate -> {
                                        if (state != null
                                                && state.contextMutable().equals(original)) {
                                            applyToContext(state, candidate.history());
                                        }
                                        return candidate.input();
                                    });
                });
    }

    /** A model view and candidate durable history. Creating it never replaces active history. */
    public record Candidate(ModelCallInput input, List<Msg> history) {
        public Candidate {
            history = List.copyOf(history);
        }
    }

    public Mono<Candidate> prepareCandidate(
            Agent agent,
            RuntimeContext ctx,
            ModelCallInput input,
            int conversationBudget,
            List<Msg> canonical) {
        return prepareCandidate(
                agent,
                ctx,
                input,
                conversationBudget,
                canonical,
                TokenCounterUtil::calculateToken,
                message -> TokenCounterUtil.calculateToken(List.of(message)));
    }

    /**
     * Prepares a candidate using the final request compiler's conversation token counter.
     * The counter excludes fixed inputs already deducted from {@code conversationBudget}.
     * The separate tool-result counter prices one result without unrelated projections or
     * fixed request overhead, using the same estimation rules.
     */
    public Mono<Candidate> prepareCandidate(
            Agent agent,
            RuntimeContext ctx,
            ModelCallInput input,
            int conversationBudget,
            List<Msg> canonical,
            ToLongFunction<List<Msg>> tokenCounter,
            ToLongFunction<Msg> toolResultTokenCounter) {
        return Mono.defer(
                () -> {
                    RuntimeContext rc = ctx == null ? RuntimeContext.empty() : ctx;
                    List<Msg> prefix = new ArrayList<>();
                    List<Msg> conversation = new ArrayList<>();
                    for (Msg msg : input.messages()) {
                        if (msg.getRole() == MsgRole.SYSTEM || msg.isRequestProjection(canonical)) {
                            prefix.add(msg);
                        } else {
                            conversation.add(msg);
                        }
                    }
                    int trigger =
                            config.getTriggerTokens() > 0
                                    ? Math.min(config.getTriggerTokens(), conversationBudget)
                                    : conversationBudget;
                    int keep =
                            config.getKeepTokens() >= 0
                                    ? config.getKeepTokens()
                                    : Math.min(
                                            config.getKeepTokensMax(),
                                            Math.max(
                                                    config.getKeepTokensMin(),
                                                    (int)
                                                            (conversationBudget
                                                                    * config
                                                                            .getKeepTokensRatio())));
                    CompactionConfig effective =
                            config.withEffective(
                                    Math.max(1, trigger),
                                    Math.min(keep, Math.max(0, conversationBudget / 2)));
                    ConversationCompactor compactor =
                            new ConversationCompactor(
                                    model,
                                    new MemoryFlushManager(workspaceManager, model),
                                    tokenCounter,
                                    toolResultTokenCounter);
                    // Copies may retain message IDs. Only an unchanged view can become durable;
                    // matching IDs alone must not persist hook-edited content or metadata.
                    boolean canCommit = sameSerializedHistory(canonical, conversation);
                    return compactor
                            .compactIfNeeded(
                                    rc,
                                    conversation,
                                    effective,
                                    agent.getName(),
                                    rc.getSessionId() == null ? "default" : rc.getSessionId())
                            .onErrorResume(
                                    error -> {
                                        if (ExceptionUtils.containsInterruptedException(error)
                                                || SessionLogException.causedBy(error))
                                            return Mono.error(error);
                                        log.warn(
                                                "Compaction failed; final request budget validation"
                                                        + " still applies",
                                                error);
                                        return Mono.just(Optional.empty());
                                    })
                            .map(
                                    result -> {
                                        if (result.isEmpty())
                                            return new Candidate(input, canonical);
                                        List<Msg> compacted = result.get();
                                        List<Msg> messages = new ArrayList<>(prefix);
                                        messages.addAll(compacted);
                                        return new Candidate(
                                                new ModelCallInput(
                                                        List.copyOf(messages),
                                                        input.tools(),
                                                        input.options(),
                                                        input.model()),
                                                canCommit ? compacted : canonical);
                                    });
                });
    }

    private static boolean sameSerializedHistory(List<Msg> canonical, List<Msg> conversation) {
        if (canonical.equals(conversation)) return true;
        if (canonical.size() != conversation.size()) return false;
        var codec = JsonUtils.getJsonCodec();
        try {
            // Normalize restored number types without rounding distinct decimal metadata values.
            return HISTORY_VALUES
                    .readTree(codec.toJson(canonical))
                    .equals(HISTORY_VALUES.readTree(codec.toJson(conversation)));
        } catch (JsonException | JsonProcessingException error) {
            // A history that cannot be serialized or parsed cannot be proven unchanged; skip
            // the writeback instead of failing the reasoning call.
            return false;
        }
    }

    private static void applyToContext(AgentState state, List<Msg> compacted) {
        if (state == null) {
            log.warn("Cannot apply compacted messages: AgentState is null");
            return;
        }
        try {
            List<Msg> ctx = state.contextMutable();
            ctx.clear();
            ctx.addAll(compacted);
            log.debug("Applied compacted messages to state ({} messages)", compacted.size());
        } catch (Exception e) {
            log.warn("Failed to apply compacted messages to state: {}", e.getMessage());
        }
    }
}