Merge pull request #4 from hrishikeshmane/fix/kiro-in-session-model-switch

fix(kiro): use session/set_model for in-session model switching

Author: hrishikeshmane

apps/server/scripts/acp-mock-agent.ts

@@ -248,6 +248,15 @@ const program = Effect.gen(function* () {
       ),
   );
 
+  yield* agent.handleSetSessionModel((request) =>
+    Effect.gen(function* () {
+      if (typeof request.modelId === "string") {
+        currentModelId = request.modelId;
+      }
+      return {};
+    }),
+  );
+
   yield* agent.handleSetSessionConfigOption((request) =>
     Effect.gen(function* () {
       if (exitOnSetConfigOption) {

apps/server/scripts/kiro-mock-agent.ts

@@ -2,7 +2,8 @@
 /**
  * Mock ACP agent that imitates kiro-cli ACP protocol for integration tests.
  *
- * Responds to: initialize, session/new, session/prompt, session/cancel.
+ * Responds to: initialize, session/new, session/prompt, session/cancel,
+ *              session/set_model, session/set_config_option.
  * Emits: session/update (agent_message_chunk), _kiro.dev/metadata extension.
  */
 import * as Effect from "effect/Effect";
@@ -14,6 +15,22 @@ import * as EffectAcpAgent from "effect-acp/agent";
 
 const sessionId = "kiro-mock-session-1";
 
+/** Available models — returned in session responses. */
+const availableModels = ["auto", "claude-sonnet-4-20250514", "claude-opus-4-20250918"];
+let currentModel = "auto";
+
+/** Build the model config option structure used in ACP responses. */
+function makeModelConfigOption(model: string) {
+  return {
+    id: "model",
+    name: "Model",
+    type: "select" as const,
+    category: "model" as const,
+    currentValue: model,
+    options: availableModels.map((m) => ({ value: m, name: m })),
+  };
+}
+
 const program = Effect.gen(function* () {
   const agent = yield* EffectAcpAgent.AcpAgent;
 
@@ -33,13 +50,32 @@ const program = Effect.gen(function* () {
   yield* agent.handleCreateSession(() =>
     Effect.succeed({
       sessionId,
+      configOptions: [makeModelConfigOption(currentModel)],
     }),
   );
 
   yield* agent.handleLoadSession(() => Effect.succeed({}));
 
   yield* agent.handleCancel(() => Effect.void);
 
+  yield* agent.handleSetSessionModel((request) =>
+    Effect.gen(function* () {
+      if (typeof request.modelId === "string") {
+        currentModel = request.modelId;
+      }
+      return {};
+    }),
+  );
+
+  yield* agent.handleSetSessionConfigOption((request) =>
+    Effect.gen(function* () {
+      if (request.configId === "model" && "value" in request && typeof request.value === "string") {
+        currentModel = request.value;
+      }
+      return { configOptions: [makeModelConfigOption(currentModel)] };
+    }),
+  );
+
   yield* agent.handlePrompt((request) =>
     Effect.gen(function* () {
       const requestedSessionId = String(request.sessionId ?? sessionId);

apps/server/src/provider/Layers/KiroAdapter.integration.test.ts

@@ -341,4 +341,141 @@ describe("KiroAdapterLive integration", () => {
         yield* adapter.stopSession(threadId);
       }).pipe(Effect.scoped, Effect.provide(adapterLayer)),
   );
+
+  it.effect(
+    "switches model in-session without restarting the process",
+    () =>
+      Effect.gen(function* () {
+        const adapter = yield* KiroAdapter;
+        const threadId = ThreadId.make("kiro-int-model-switch-1");
+
+        yield* adapter.startSession({
+          threadId,
+          provider: "kiro",
+          cwd: process.cwd(),
+          runtimeMode: "full-access",
+          modelSelection: { provider: "kiro", model: "auto" },
+        });
+
+        const spawnCountAfterStart = capturedArgs.length;
+
+        // First turn with default model
+        yield* adapter.sendTurn({
+          threadId,
+          input: "first turn",
+          attachments: [],
+        });
+
+        // Second turn with different model — should NOT respawn
+        yield* adapter.sendTurn({
+          threadId,
+          input: "second turn after model switch",
+          attachments: [],
+          modelSelection: {
+            provider: "kiro",
+            model: "claude-sonnet-4-20250514",
+          },
+        });
+
+        // Session should NOT have been restarted — only one spawn
+        expect(capturedArgs.length).toBe(spawnCountAfterStart);
+
+        yield* adapter.stopSession(threadId);
+      }).pipe(Effect.scoped, Effect.provide(adapterLayer)),
+  );
+
+  it.effect(
+    "updates session.model immediately after in-session model switch",
+    () =>
+      Effect.gen(function* () {
+        const adapter = yield* KiroAdapter;
+        const threadId = ThreadId.make("kiro-int-model-state-1");
+
+        const session = yield* adapter.startSession({
+          threadId,
+          provider: "kiro",
+          cwd: process.cwd(),
+          runtimeMode: "full-access",
+          modelSelection: { provider: "kiro", model: "auto" },
+        });
+
+        expect(session.model).toBe("auto");
+
+        // Collect turn.started events to verify model is correct
+        const eventsFiber = yield* adapter.streamEvents.pipe(
+          Stream.filter(
+            (event) => event.type === "turn.started" || event.type === "turn.completed",
+          ),
+          Stream.take(2), // turn.started + turn.completed for the model-switch turn
+          Stream.runCollect,
+          Effect.forkChild,
+        );
+
+        // Send turn with new model
+        const turn = yield* adapter.sendTurn({
+          threadId,
+          input: "switch model turn",
+          attachments: [],
+          modelSelection: {
+            provider: "kiro",
+            model: "claude-sonnet-4-20250514",
+          },
+        });
+
+        expect(turn.threadId).toBe(threadId);
+
+        const events = yield* Fiber.join(eventsFiber);
+        const turnStarted = events.find((e) => e.type === "turn.started");
+        expect(turnStarted).toBeDefined();
+        if (turnStarted?.type === "turn.started") {
+          expect(turnStarted.payload.model).toBe("claude-sonnet-4-20250514");
+        }
+
+        // Verify session state reflects the new model
+        const sessions = yield* adapter.listSessions();
+        const currentSession = sessions.find((s) => s.threadId === threadId);
+        expect(currentSession?.model).toBe("claude-sonnet-4-20250514");
+
+        yield* adapter.stopSession(threadId);
+      }).pipe(Effect.scoped, Effect.provide(adapterLayer)),
+  );
+
+  it.effect(
+    "does not call setModel when model is unchanged between turns",
+    () =>
+      Effect.gen(function* () {
+        const adapter = yield* KiroAdapter;
+        const threadId = ThreadId.make("kiro-int-model-unchanged-1");
+
+        yield* adapter.startSession({
+          threadId,
+          provider: "kiro",
+          cwd: process.cwd(),
+          runtimeMode: "full-access",
+          modelSelection: { provider: "kiro", model: "auto" },
+        });
+
+        // Two turns with the same model — should not trigger setModel
+        yield* adapter.sendTurn({
+          threadId,
+          input: "first turn",
+          attachments: [],
+          modelSelection: { provider: "kiro", model: "auto" },
+        });
+
+        yield* adapter.sendTurn({
+          threadId,
+          input: "second turn same model",
+          attachments: [],
+          modelSelection: { provider: "kiro", model: "auto" },
+        });
+
+        // Both turns should succeed without issues (no setModel called)
+        const sessions = yield* adapter.listSessions();
+        const currentSession = sessions.find((s) => s.threadId === threadId);
+        expect(currentSession?.model).toBe("auto");
+
+        yield* adapter.stopSession(threadId);
+      }).pipe(Effect.scoped, Effect.provide(adapterLayer)),
+  );
 });

apps/server/src/provider/Layers/KiroAdapter.ts

@@ -984,13 +984,26 @@ function makeKiroAdapter(options?: KiroAdapterLiveOptions) {
 
         // Only switch model if different from current
         if (model !== ctx.session.model && model !== "auto") {
-          yield* ctx.acp.setModel(model).pipe(
-            Effect.mapError((error) =>
-              mapAcpToAdapterError(PROVIDER, input.threadId, "session/set_model", error),
-            ),
-            // Kiro may not support session/set_config_option — ignore errors
-            Effect.catch(() => Effect.void),
-          );
+          // Kiro supports session/set_model but NOT session/set_config_option.
+          // AcpSessionRuntime.setModel routes through setConfigOption, which
+          // Kiro rejects with -32601 "Method not found". Call the correct
+          // RPC method directly via the raw request interface.
+          yield* ctx.acp
+            .request("session/set_model", {
+              sessionId: ctx.mainSessionId,
+              modelId: model,
+            })
+            .pipe(
+              Effect.mapError((error) =>
+                mapAcpToAdapterError(PROVIDER, input.threadId, "session/set_model", error),
+              ),
+            );
+          // Update session model immediately after successful setModel,
+          // matching the pattern used by ClaudeAdapter and OpenCodeAdapter.
+          ctx.session = {
+            ...ctx.session,
+            model,
+          };
         }
         ctx.activeTurnId = turnId;
         ctx.session = {
@@ -1036,7 +1049,6 @@ function makeKiroAdapter(options?: KiroAdapterLiveOptions) {
           ...ctx.session,
           activeTurnId: turnId,
           updatedAt: yield* nowIso,
-          model,
         };
 
         const stopReason = ctx.interrupted ? "interrupted" : (result.stopReason ?? null);