diff --git a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java index 0bb6b49af7..c5c84713fd 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java +++ b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java @@ -2118,7 +2118,7 @@ Flux reasoningStream( rc, MiddlewareBase::onModelCall, modelCallCore) - .apply(new ModelCallInput(messages, tools, options, modelForCall())) + .apply(new ModelCallInput(messages, tools, options, modelForCall(rc))) .doOnNext(this::publishEvent); } @@ -3052,7 +3052,7 @@ protected Mono summarizing() { List messageList = prepareSummaryMessages(); GenerateOptions generateOptions = buildGenerateOptions(); ReasoningContext context = new ReasoningContext(getName()); - Model summaryModel = modelForCall(); + Model summaryModel = modelForCall(rc); publishEvent(new ExceedMaxItersEvent("", maxIters, maxIters)); return hookDispatcher @@ -3482,29 +3482,63 @@ protected GenerateOptions buildGenerateOptions() { return baseOptions != null ? baseOptions : GenerateOptions.builder().build(); } - private Model modelForCall() { - Model fallbackModel = modelConfig.fallbackModel(); - if (fallbackModel == null) { - return model; + /** + * Resolves the model for a single call, wrapping it with the configured fallback when one is + * present. + */ + private Model modelForCall(RuntimeContext ctx) { + Model primary = getModel(ctx); + Model fallback = modelConfig.fallbackModel(); + return fallback == null ? primary : wrapWithFallback(primary, fallback); + } + + /** + * Returns the model effective for the session carried by {@code ctx}: the session's registered + * id resolved through {@link ModelRegistry} when set, otherwise the agent's default model. An + * unresolvable session id is logged and falls back to the default. + * + * @param ctx the runtime context identifying the session (may be {@code null}) + * @return the resolved session model, or the default model when no id is set or it fails to + * resolve + */ + public Model getModel(RuntimeContext ctx) { + AgentState s = ctx != null ? ctx.getAgentState() : null; + String mid = s != null ? s.getModelId() : null; + if (mid != null && !mid.isBlank()) { + try { + return ModelRegistry.resolve(mid); + } catch (Exception e) { + log.warn( + "Session modelId [{}] unresolvable, fallback to default model: {}", + mid, + e.getMessage()); + } } + return model; + } - AtomicReference activeModel = new AtomicReference<>(model); + /** + * Wraps {@code primary} so that {@code fallback} takes over if the first signal from + * {@code primary} is an error. + */ + private static Model wrapWithFallback(Model primary, Model fallback) { + AtomicReference activeModel = new AtomicReference<>(primary); return new Model() { @Override public Flux stream( List messages, List tools, GenerateOptions options) { - Flux primaryFlux = model.stream(messages, tools, options); + Flux primaryFlux = primary.stream(messages, tools, options); return primaryFlux.switchOnFirst( (signal, flux) -> { if (signal.isOnError()) { Throwable error = signal.getThrowable(); - activeModel.set(fallbackModel); + activeModel.set(fallback); log.warn( "Primary model {} failed, switching to fallback {}", - model.getModelName(), - fallbackModel.getModelName(), + primary.getModelName(), + fallback.getModelName(), error); - return fallbackModel.stream(messages, tools, options); + return fallback.stream(messages, tools, options); } return flux; }); @@ -3693,6 +3727,54 @@ public void setPermissionMode(RuntimeContext ctx, PermissionMode mode) { setPermissionMode(uid, sid, mode); } + /** + * Switches the model used by the given {@code (userId, sessionId)} session at runtime by storing + * its id in the persisted {@link AgentState}. The id is resolved through {@link ModelRegistry} + * on each {@code call}, so a blank or {@code null} id restores the agent's default model. The + * change is persisted so the next {@code call} on that session—and any cluster node that + * reloads the state—sees it. + * + *

Unlike {@link #setPermissionMode(String, String, PermissionMode)}, no per-slot cache is + * rebuilt: models are stateless, so resolution happens fresh on every call. + * + * @param userId user identity for the slot (may be {@code null}) + * @param sessionId session identity (falls back to the default session id when {@code null}) + * @param modelId the registered model id to use, or {@code null}/blank to restore the default + */ + public void setModel(String userId, String sessionId, String modelId) { + String sid = (sessionId == null || sessionId.isBlank()) ? defaultSessionId : sessionId; + AgentState s = getAgentState(userId, sid); + s.setModelId(modelId); + saveAgentState(userId, sid); + } + + /** + * Switches the model for the session identified by the given {@link RuntimeContext}. See + * {@link #setModel(String, String, String)}. + * + * @param ctx the runtime context identifying the session + * @param modelId the registered model id to use, or {@code null}/blank to restore the default + */ + public void setModel(RuntimeContext ctx, String modelId) { + String uid = ctx != null ? ctx.getUserId() : null; + String sid = ctx != null ? ctx.getSessionId() : null; + setModel(uid, sid, modelId); + } + + /** + * Clears the session's model id so subsequent calls revert to the agent's default model. The + * change is persisted. + * + * @param userId user identity for the slot (may be {@code null}) + * @param sessionId session identity (falls back to the default session id when {@code null}) + */ + public void clearModel(String userId, String sessionId) { + String sid = (sessionId == null || sessionId.isBlank()) ? defaultSessionId : sessionId; + AgentState s = getAgentState(userId, sid); + s.setModelId(null); + saveAgentState(userId, sid); + } + /** * Returns the current {@link PermissionMode} for the given {@code (userId, sessionId)} session. * diff --git a/agentscope-core/src/main/java/io/agentscope/core/state/AgentState.java b/agentscope-core/src/main/java/io/agentscope/core/state/AgentState.java index f1dc6e7d8f..5a69bb182c 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/state/AgentState.java +++ b/agentscope-core/src/main/java/io/agentscope/core/state/AgentState.java @@ -54,7 +54,8 @@ "permission_context", "tool_context", "tasks_context", - "plan_mode_context" + "plan_mode_context", + "model_id" }) public final class AgentState implements State { @@ -69,6 +70,7 @@ public final class AgentState implements State { private final ToolContextState toolContext; private final TaskContextState tasksContext; private final PlanModeContextState planModeContext; + private String modelId; /** * Per-session interrupt signal. Runtime-only (never serialized): a stateless agent engine @@ -100,6 +102,7 @@ private AgentState(Builder builder) { builder.planModeContext == null ? new PlanModeContextState() : builder.planModeContext; + this.modelId = builder.modelId; } @JsonCreator @@ -114,7 +117,8 @@ static AgentState fromJson( @JsonProperty("permission_context") PermissionContextState permissionContext, @JsonProperty("tool_context") ToolContextState toolContext, @JsonProperty("tasks_context") TaskContextState tasksContext, - @JsonProperty("plan_mode_context") PlanModeContextState planModeContext) { + @JsonProperty("plan_mode_context") PlanModeContextState planModeContext, + @JsonProperty("model_id") String modelId) { Builder b = builder(); if (sessionId != null) { b.sessionId(sessionId); @@ -149,6 +153,9 @@ static AgentState fromJson( if (planModeContext != null) { b.planModeContext(planModeContext); } + if (modelId != null) { + b.modelId(modelId); + } return b.build(); } @@ -247,6 +254,15 @@ public PlanModeContextState getPlanModeContext() { return planModeContext; } + @JsonProperty("model_id") + public String getModelId() { + return modelId; + } + + public void setModelId(String modelId) { + this.modelId = (modelId == null || modelId.isBlank()) ? null : modelId; + } + /** * The per-session interrupt signal for this state, created on first access. Runtime-only and not * serialized. @@ -305,7 +321,8 @@ public boolean equals(Object o) { && Objects.equals(permissionContext, other.permissionContext) && Objects.equals(toolContext, other.toolContext) && Objects.equals(tasksContext, other.tasksContext) - && Objects.equals(planModeContext, other.planModeContext); + && Objects.equals(planModeContext, other.planModeContext) + && Objects.equals(modelId, other.modelId); } @Override @@ -321,7 +338,8 @@ public int hashCode() { permissionContext, toolContext, tasksContext, - planModeContext); + planModeContext, + modelId); } @Override @@ -351,6 +369,7 @@ public static final class Builder { private ToolContextState toolContext; private TaskContextState tasksContext; private PlanModeContextState planModeContext; + private String modelId; private Builder() {} @@ -416,6 +435,11 @@ public Builder planModeContext(PlanModeContextState planModeContext) { return this; } + public Builder modelId(String modelId) { + this.modelId = modelId; + return this; + } + public AgentState build() { return new AgentState(this); } diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java index 5d88b1febc..4186b18abf 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java @@ -18,6 +18,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -32,6 +33,7 @@ import io.agentscope.core.model.ChatModelBase; import io.agentscope.core.model.ChatResponse; import io.agentscope.core.model.GenerateOptions; +import io.agentscope.core.model.ModelRegistry; import io.agentscope.core.model.ToolSchema; import io.agentscope.core.state.AgentState; import io.agentscope.core.state.InMemoryAgentStateStore; @@ -41,6 +43,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import java.util.stream.IntStream; import org.junit.jupiter.api.DisplayName; @@ -69,6 +72,40 @@ protected Flux doStream( } } + private static final class NamedModel extends ChatModelBase { + private final String name; + private final AtomicInteger streamCount = new AtomicInteger(); + + private NamedModel(String name) { + this.name = name; + } + + @Override + public String getModelName() { + return name; + } + + int streamCount() { + return streamCount.get(); + } + + @Override + protected Flux doStream( + List messages, List tools, GenerateOptions options) { + streamCount.incrementAndGet(); + return Flux.just( + ChatResponse.builder() + .content(List.of(TextBlock.builder().text("ok").build())) + .build()); + } + } + + private RuntimeContext ctxWithState(ReActAgent agent, String uid, String sid) { + RuntimeContext ctx = RuntimeContext.builder().userId(uid).sessionId(sid).build(); + ctx.setAgentState(agent.getAgentState(uid, sid)); + return ctx; + } + private ReActAgent agent(InMemoryAgentStateStore store) { return ReActAgent.builder() .name("asst") @@ -307,4 +344,101 @@ void concurrentStreamEventsAreIsolated() { assertEquals(1, ends, "each stream must be closed by exactly one AgentEndEvent"); } } + + @Test + @DisplayName("setModel persists modelId per session and resolves via ModelRegistry") + void setModelPersistsAndResolves() { + InMemoryAgentStateStore store = new InMemoryAgentStateStore(); + ReActAgent agent = agent(store); + ModelRegistry.register("test:model-a", new NamedModel("model-a")); + try { + agent.setModel("u1", "sessA", "test:model-a"); + assertEquals("test:model-a", agent.getAgentState("u1", "sessA").getModelId()); + + assertEquals( + "model-a", agent.getModel(ctxWithState(agent, "u1", "sessA")).getModelName()); + + // persisted across engines (cluster drift: fresh node loads from the same store) + ReActAgent reborn = agent(store); + assertEquals("test:model-a", reborn.getAgentState("u1", "sessA").getModelId()); + assertEquals( + "model-a", reborn.getModel(ctxWithState(reborn, "u1", "sessA")).getModelName()); + } finally { + ModelRegistry.reset(); + } + } + + @Test + @DisplayName("clearModel reverts to the default model") + void clearModelRevertsToDefault() { + ReActAgent agent = agent(new InMemoryAgentStateStore()); + ModelRegistry.register("test:model-b", new NamedModel("model-b")); + try { + agent.setModel("u1", "sessA", "test:model-b"); + assertEquals( + "model-b", agent.getModel(ctxWithState(agent, "u1", "sessA")).getModelName()); + + agent.clearModel("u1", "sessA"); + assertNull(agent.getAgentState("u1", "sessA").getModelId()); + assertEquals("noop", agent.getModel(ctxWithState(agent, "u1", "sessA")).getModelName()); + } finally { + ModelRegistry.reset(); + } + } + + @Test + @DisplayName("unresolvable session modelId falls back to default without throwing") + void unresolvableModelIdFallsBack() { + ReActAgent agent = agent(new InMemoryAgentStateStore()); + agent.setModel("u1", "sessA", "test:does-not-exist"); + assertEquals("noop", agent.getModel(ctxWithState(agent, "u1", "sessA")).getModelName()); + } + + @Test + @DisplayName("per-session modelId is isolated across sessions") + void modelIdIsolatedPerSession() { + ReActAgent agent = agent(new InMemoryAgentStateStore()); + ModelRegistry.register("test:model-a", new NamedModel("model-a")); + ModelRegistry.register("test:model-b", new NamedModel("model-b")); + try { + agent.setModel("u1", "sessA", "test:model-a"); + agent.setModel("u1", "sessB", "test:model-b"); + + assertEquals( + "model-a", agent.getModel(ctxWithState(agent, "u1", "sessA")).getModelName()); + assertEquals( + "model-b", agent.getModel(ctxWithState(agent, "u1", "sessB")).getModelName()); + } finally { + ModelRegistry.reset(); + } + } + + @Test + @DisplayName("call() routes reasoning through the per-session model") + void callRoutesThroughPerSessionModel() { + InMemoryAgentStateStore store = new InMemoryAgentStateStore(); + NamedModel defaultModel = new NamedModel("default"); + NamedModel sessionModel = new NamedModel("session"); + ReActAgent agent = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(defaultModel) + .stateStore(store) + .build(); + ModelRegistry.register("test:session-model", sessionModel); + try { + agent.setModel("u1", "sessA", "test:session-model"); + RuntimeContext ctx = RuntimeContext.builder().userId("u1").sessionId("sessA").build(); + agent.call(List.of(userMsg("hello")), ctx).block(Duration.ofSeconds(10)); + + assertTrue( + sessionModel.streamCount() > 0, + "reasoning must route through the per-session model"); + assertEquals( + 0, defaultModel.streamCount(), "default model must not be used after setModel"); + } finally { + ModelRegistry.reset(); + } + } } diff --git a/agentscope-core/src/test/java/io/agentscope/core/state/AgentStateTest.java b/agentscope-core/src/test/java/io/agentscope/core/state/AgentStateTest.java index 54a30c8f1b..aa8062135a 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/state/AgentStateTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/state/AgentStateTest.java @@ -17,6 +17,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -165,4 +166,32 @@ void jsonOmittedFieldsFallBackToDefaults() throws Exception { assertEquals(0, decoded.getCurIter()); assertNotNull(decoded.getReplyId()); } + + @Test + void modelIdRoundTrip() throws Exception { + AgentState original = AgentState.builder().sessionId("s").modelId("deepseek:v4").build(); + String json = mapper.writeValueAsString(original); + assertTrue(json.contains("\"model_id\":\"deepseek:v4\""), () -> json); + AgentState decoded = mapper.readValue(json, AgentState.class); + assertEquals("deepseek:v4", decoded.getModelId()); + } + + @Test + void modelIdOmittedInOldJsonFallsBackToNull() throws Exception { + AgentState decoded = mapper.readValue("{\"session_id\":\"only-id\"}", AgentState.class); + assertNull(decoded.getModelId()); + } + + @Test + void setModelIdNormalizesBlankToNull() { + AgentState s = AgentState.builder().build(); + s.setModelId("deepseek:v4"); + assertEquals("deepseek:v4", s.getModelId()); + s.setModelId(""); + assertNull(s.getModelId()); + s.setModelId(" "); + assertNull(s.getModelId()); + s.setModelId(null); + assertNull(s.getModelId()); + } } diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java index d669184cb9..3ef06c8bd1 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java @@ -372,6 +372,34 @@ public io.agentscope.core.permission.PermissionMode getPermissionMode( return delegate.getPermissionMode(userId, sessionId); } + /** + * Switches the model for the session identified by the given {@link RuntimeContext}, persisted + * in the session's {@link io.agentscope.core.state.AgentState}. + */ + public void setModel(RuntimeContext ctx, String modelId) { + delegate.setModel(ctx, modelId); + } + + /** + * Switches the model for the given {@code (userId, sessionId)} session at runtime by storing its + * id in the persisted session state. + */ + public void setModel(String userId, String sessionId, String modelId) { + delegate.setModel(userId, sessionId, modelId); + } + + /** @return the model effective for the session carried by {@code ctx}. */ + public Model getModel(RuntimeContext ctx) { + return delegate.getModel(ctx); + } + + /** + * Clears the session's model id so subsequent calls revert to the default model. + */ + public void clearModel(String userId, String sessionId) { + delegate.clearModel(userId, sessionId); + } + @Override public void close() { try { @@ -2140,7 +2168,7 @@ public HarnessAgent build() { inner.middleware( new MemoryFlushMiddleware( wsManager, - memoryModel, + memoryConfig.model(), effectiveFlushPrompt, memoryConfig.flushTrigger(), effectiveIsolationScope)); @@ -2166,13 +2194,8 @@ public HarnessAgent build() { } CompactionMiddleware compactionHook = null; if (!disableCompaction && compactionConfig != null) { - Model compactionModel = - compactionConfig.getModel() != null ? compactionConfig.getModel() : model; - if (compactionModel != null) { - compactionHook = - new CompactionMiddleware(wsManager, compactionModel, compactionConfig); - inner.middleware(compactionHook); - } + compactionHook = new CompactionMiddleware(wsManager, compactionConfig); + inner.middleware(compactionHook); } if (!disableToolResultEviction && toolResultEvictionConfig != null) { inner.middleware( diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/CompactionMiddleware.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/CompactionMiddleware.java index ae6405be3d..3750fc1ea3 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/CompactionMiddleware.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/CompactionMiddleware.java @@ -59,13 +59,10 @@ public class CompactionMiddleware implements HarnessRuntimeMiddleware { private static final Logger log = LoggerFactory.getLogger(CompactionMiddleware.class); private final WorkspaceManager workspaceManager; - private final Model model; private final CompactionConfig config; - public CompactionMiddleware( - WorkspaceManager workspaceManager, Model model, CompactionConfig config) { + public CompactionMiddleware(WorkspaceManager workspaceManager, CompactionConfig config) { this.workspaceManager = workspaceManager; - this.model = model; this.config = config; } @@ -79,6 +76,8 @@ public Flux onReasoning( return next.apply(input); } final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty(); + final Model dedicated = config.getModel(); + final Model effective = dedicated != null ? dedicated : reActAgent.getModel(rc); return Flux.defer( () -> { @@ -98,12 +97,12 @@ public Flux onReasoning( String sessionId = rc != null && rc.getSessionId() != null ? rc.getSessionId() : "default"; - CompactionConfig effectiveConfig = resolveEffectiveConfig(); + CompactionConfig effectiveConfig = resolveEffectiveConfig(effective); MemoryFlushManager flushManager = - new MemoryFlushManager(workspaceManager, model); + new MemoryFlushManager(workspaceManager, effective); ConversationCompactor compactor = - new ConversationCompactor(model, flushManager); + new ConversationCompactor(effective, flushManager); final Msg sys = systemMsg; return compactor @@ -145,7 +144,7 @@ public Flux onReasoning( /** * Resolves dynamic defaults in the config using the model's context window. */ - private CompactionConfig resolveEffectiveConfig() { + private CompactionConfig resolveEffectiveConfig(Model effective) { int configTrigger = config.getTriggerTokens(); int configKeep = config.getKeepTokens(); @@ -154,7 +153,7 @@ private CompactionConfig resolveEffectiveConfig() { return config; } - int contextWindow = model.getContextWindowSize(); + int contextWindow = effective.getContextWindowSize(); int effectiveTrigger; if (configTrigger == 0) { diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java index 1ea976e54a..499154348f 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java @@ -15,6 +15,7 @@ */ package io.agentscope.harness.agent.middleware; +import io.agentscope.core.ReActAgent; import io.agentscope.core.agent.Agent; import io.agentscope.core.agent.RuntimeContext; import io.agentscope.core.event.AgentEvent; @@ -164,8 +165,16 @@ private Mono doFlush(Agent agent, RuntimeContext rc) { return Mono.empty(); } + Model effective; + if (model != null) { + effective = model; + } else if (agent instanceof ReActAgent reActAgent) { + effective = reActAgent.getModel(rc); + } else { + return Mono.empty(); + } MemoryFlushManager flushManager = - new MemoryFlushManager(workspaceManager, model, flushPrompt); + new MemoryFlushManager(workspaceManager, effective, flushPrompt); boolean shouldFlush = shouldFlushNow(rc); Mono flushMono;