spring-projects / spring-projects/spring-ai

Bug: DefaultAroundAdvisorChain fails to maintain immutability across multiple calls

Open
#3,537 2 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

status: waiting-for-triage
Dominant language
Java
Stars
9.5k
Forks
2.9k
Avg merge
1d 7h
Merged PRs (30d)
6

Description

Bug description
The current implementation of DefaultAroundAdvisorChain has a design issue where the chain's state is being polluted after the first call, making it impossible to reuse the chain for multiple calls. This violates the immutability principle that should be maintained by the chain.

Environment

  • Spring AI version: 1.0.0
  • Java version: ALL
  • No vector store involved

Steps to reproduce

  1. Create a chain with multiple advisors (including intermediate and terminal advisors)
  2. Attempt to call the same chain instance multiple times
  3. Observe that the chain's state is being polluted after the first call

Expected behavior

  • The chain should be truly immutable
  • Multiple calls to the same chain instance should produce consistent results
  • The chain should be reusable without any state pollution
  • Each call should maintain the correct execution order: inter1 before -> inter2 before -> terminal -> inter2 after -> inter1 after

Minimal Complete Reproducible example

package org.springframework.ai.chat.client.advisor;

import io.micrometer.observation.ObservationRegistry;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.prompt.Prompt;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/*
 * This is a real test that uses the DefaultAroundAdvisorChain to execute a chain of advisors.
 * We create a terminal advisor and two intermediate advisors.
 * A stack in the context is used to track the execution order of the advisors.
 * It should be: inter1 before -> inter2 before -> terminal -> inter2 after -> inter1 after
 */
public class DefaultAroundAdvisorChainRealTests {
    static final String stackName = "testStack";

    static class TerminalCallAdvisor implements CallAdvisor {

        @Override
        public int getOrder() {
            return 2147483647;
        }

        @Override
        public String getName() {
            return "terminal-" + getOrder();
        }

        @Override
        public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
            List<String> stack = (List<String>) request.context().get(stackName);
            stack.add(this.getName() + " executed");
            return ChatClientResponse.builder()
                    .context(stackName, stack)
                    .build();
        }
    }


    static class IntermediateCallAdvisor implements CallAdvisor {
        private final int order;

        IntermediateCallAdvisor(int order) {
            this.order = order;
        }

        @Override
        public int getOrder() {
            return order;
        }

        @Override
        public String getName() {
            return "intermediate-" + order;
        }

        @Override
        public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
            List<String> stack = (List<String>) request.context().get(stackName);
            stack.add(getName() + " before execution");

            var response = chain.nextCall(request);
            // updatedStack == stack
            List<String> updatedStack = (List<String>) response.context().get(stackName);
            updatedStack.add(getName() + " after execution");
            return response;
        }
    }

    @Test
    void realTest() {

        List<CallAdvisor> advisors = new ArrayList<>(List.of(
                new IntermediateCallAdvisor(100),
                new IntermediateCallAdvisor(200),
                new TerminalCallAdvisor()
        ));
        // Chain will re-order advisors, so we can shuffle advisors before passing them to chain.
        Collections.shuffle(advisors);

        CallAdvisorChain chain = DefaultAroundAdvisorChain.builder(ObservationRegistry.NOOP)
                .pushAll(advisors)
                .build();

        //  AdvisorChain is immutable and stateless, so we can call it multiple times
        for (int i = 0; i < 100; i++) {

            ChatClientRequest request = ChatClientRequest.builder()
                    .prompt(new Prompt("Hello"))
                    .context(Map.of(stackName, new ArrayList<>()))
                    .build();

            ChatClientResponse response = chain.nextCall(request);
            List<String> stack = (List<String>) response.context().get(stackName);
            Assertions.assertNotNull(stack);
            Assertions.assertEquals(5, stack.size());
            List<String> expectedStack = List.of(
                    "intermediate-100 before execution",
                    "intermediate-200 before execution",
                    "terminal-2147483647 executed",
                    "intermediate-200 after execution",
                    "intermediate-100 after execution"
            );
            Assertions.assertEquals(expectedStack, stack);
        }
    }
}

Proposed Fix
I've implemented a simpler and more correct version of the chain that properly maintains immutability. Here's my implementation:

package org.springframework.ai.chat.client.advisor;

import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.api.BaseAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationContext;
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationConvention;
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationDocumentation;
import org.springframework.ai.chat.client.advisor.observation.DefaultAdvisorObservationConvention;
import org.springframework.ai.template.TemplateRenderer;
import org.springframework.ai.template.st.StTemplateRenderer;
import org.springframework.core.OrderComparator;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Objects;

public class DefaultAroundAdvisorChain implements BaseAdvisorChain {

    public static final AdvisorObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultAdvisorObservationConvention();

    private static final TemplateRenderer DEFAULT_TEMPLATE_RENDERER = StTemplateRenderer.builder().build();

    private final ObservationRegistry observationRegistry;
    private final TemplateRenderer templateRenderer;

    private final List<CallAdvisor> callAdvisors;
    private final List<StreamAdvisor> streamAdvisors;

    public DefaultAroundAdvisorChain(ObservationRegistry observationRegistry, @Nullable TemplateRenderer templateRenderer,
                                      Collection<CallAdvisor> callAdvisors, Collection<StreamAdvisor> streamAdvisors) {
        Assert.notNull(observationRegistry, "the observationRegistry must be non-null");
        Assert.notNull(callAdvisors, "the callAdvisors must be non-null");
        Assert.notNull(streamAdvisors, "the streamAdvisors must be non-null");

        this.observationRegistry = observationRegistry;
        this.templateRenderer = templateRenderer == null ? DEFAULT_TEMPLATE_RENDERER : templateRenderer;
        this.callAdvisors = List.copyOf(callAdvisors);
        this.streamAdvisors = List.copyOf(streamAdvisors);
    }

    public static Builder builder(ObservationRegistry observationRegistry) {
        return new Builder(observationRegistry);
    }

    @Override
	public ChatClientResponse nextCall(ChatClientRequest request) {
		Assert.notNull(request, "the chatClientRequest cannot be null");

        if (this.callAdvisors.isEmpty()) {
			throw new IllegalStateException("No CallAdvisors available to execute");
        }

        var advisor = this.callAdvisors.get(0);
        var nextChain = this.mutate()
                .callAdvisors(this.callAdvisors.subList(1, this.callAdvisors.size()))
                .build();

        var observationBuilder = AdvisorObservationContext.builder()
                .advisorName(advisor.getName())
                .chatClientRequest(request)
                .order(advisor.getOrder());

        var response = AdvisorObservationDocumentation.AI_ADVISOR
                .observation(null, DEFAULT_OBSERVATION_CONVENTION, observationBuilder::build, this.observationRegistry)
                .observe(() -> advisor.adviseCall(request, nextChain));
        return Objects.requireNonNull(response);
    }

    @Override
    public Flux<ChatClientResponse> nextStream(@NonNull ChatClientRequest request) {
        Assert.notNull(request, "the request must be non-null");

        return Flux.deferContextual(contextView -> {
            if (this.streamAdvisors.isEmpty()) {
                return Flux.error(new IllegalStateException("No StreamAdvisors available"));
            }

            var advisor = this.streamAdvisors.get(0);
            var nextChain = this.mutate()
                    .streamAdvisors(this.streamAdvisors.subList(1, this.streamAdvisors.size()))
                    .build();

            var observationBuilder = AdvisorObservationContext.builder()
                    .advisorName(advisor.getName())
                    .chatClientRequest(request)
                    .order(advisor.getOrder());

            var observation = AdvisorObservationDocumentation.AI_ADVISOR.observation(null,
                    DEFAULT_OBSERVATION_CONVENTION, observationBuilder::build, this.observationRegistry);

            observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();

            return Flux.defer(() -> advisor.adviseStream(request, nextChain)
                    .doOnError(observation::error)
                    .doFinally(s -> observation.stop())
                    .contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation)));
        });
    }

    public Builder mutate() {
        return new Builder(this.observationRegistry)
                .callAdvisors(this.callAdvisors)
                .streamAdvisors(this.streamAdvisors)
                .templateRenderer(this.templateRenderer);
    }

    @Override
    public List<CallAdvisor> getCallAdvisors() {
        return this.callAdvisors;
    }

    @Override
    public List<StreamAdvisor> getStreamAdvisors() {
        return this.streamAdvisors;
    }

    @Override
    public ObservationRegistry getObservationRegistry() {
        return this.observationRegistry;
    }

    public static class Builder {

        private final ObservationRegistry observationRegistry;

        private final List<CallAdvisor> callAdvisors;

        private final List<StreamAdvisor> streamAdvisors;

        private TemplateRenderer templateRenderer;

        public Builder(ObservationRegistry observationRegistry) {
            this.observationRegistry = observationRegistry;
            this.callAdvisors = new ArrayList<>();
            this.streamAdvisors = new ArrayList<>();
        }

        public Builder templateRenderer(TemplateRenderer templateRenderer) {
            this.templateRenderer = templateRenderer;

            return this;
        }

        public Builder callAdvisors(Collection<CallAdvisor> callAdvisors) {
            this.callAdvisors.clear();
            this.callAdvisors.addAll(callAdvisors);
            OrderComparator.sort(this.callAdvisors);

            return this;
        }

        public Builder streamAdvisors(Collection<StreamAdvisor> streamAdvisors) {
            this.streamAdvisors.clear();
            this.streamAdvisors.addAll(streamAdvisors);
            OrderComparator.sort(this.streamAdvisors);

            return this;
        }

        /**
         * @deprecated bad idea/name.
         */
        @Deprecated()
        public Builder push(Advisor advisor) {
            Assert.notNull(advisor, "the advisor must be non-null");

            return this.pushAll(List.of(advisor));
        }

        /**
         * @deprecated bad idea/name. Use {@link #setAdvisors(List)} instead.
         */
        @Deprecated()
        public Builder pushAll(List<? extends Advisor> advisors) {
            return this.setAdvisors(advisors);
        }

        /**
         *
         * @see #streamAdvisors(Collection)
         * @see #callAdvisors(Collection)
         */
        public Builder setAdvisors(List<? extends Advisor> advisors) {
            Assert.notNull(advisors, "the advisors must be non-null");
            Assert.noNullElements(advisors, "the advisors must not contain null elements");

            List<CallAdvisor> callAdvisors = advisors.stream()
                    .filter(CallAdvisor.class::isInstance)
                    .map(CallAdvisor.class::cast)
                    .toList();
            this.callAdvisors(callAdvisors);

            List<StreamAdvisor> streamAdvisors = advisors.stream()
                    .filter(StreamAdvisor.class::isInstance)
                    .map(StreamAdvisor.class::cast)
                    .toList();
            this.streamAdvisors(streamAdvisors);

            return this;
        }

        public DefaultAroundAdvisorChain build() {
            return new DefaultAroundAdvisorChain(
                    this.observationRegistry, this.templateRenderer,
                    this.callAdvisors, this.streamAdvisors);
        }
    }
}

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at DefaultAroundAdvisorChain.nextCall and run the supplied realTest reproduction with the intermediate and terminal advisors. Check repeated calls against the expected five-entry execution order; done means the same chain can be called repeatedly without state pollution and each response has the expected order.

Written by the indexing model from the issue text.

Assessment

Tech stack
java, spring
Domain
backend
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.