Skip to content

Commit f0fa43b

Browse files
author
hrishi mane
committed
fix(kiro): in-session model + agent switching via set_model/set_mode RPC
Supersedes PR #4's set_model work with a corrected implementation covering both models (`session/set_model`) and agents (`session/set_mode`), and closes the first-turn alignment gotcha PR #4 had. Key learnings from debugging sessions on threads e3a5c1bf (ValidationException), 55afbe27 (silent model swap no-op), 23f518fe (session/load ignores --model), 3e377f8b (UI not rendering post-respawn reply), and 593e16ae (final working state with 3 set_mode + 4 set_model + 7 prompts, zero respawns): - Kiro's `session/set_model` RPC only works if kiro-cli acp was spawned WITHOUT the `--model` flag. Passing `--model <slug>` at spawn locks the model and makes later set_model calls silently no-op. - Kiro's `session/set_mode` RPC is the agent-equivalent of set_model. Kirodex uses it (connection.setSessionMode) and it works. Agent switching no longer needs to respawn kiro-cli, eliminating the UI artifact where post-respawn replies didn't render cleanly. - effect-acp's agent SDK doesn't wire `session/set_mode` through its core handler dispatcher. `handleExtRequest("session/set_mode", ...)` is the correct hook on the mock side. - AcpSessionRuntime.setConfigOption routes through `session/set_config_option` which Kiro rejects with -32601. We bypass the upstream helper and issue set_model / set_mode via the raw request channel (`ctx.acp.request(method, payload)`) — same pattern as PR #4. - `ctx.activeMode` and `ctx.activeModel` track the agent/model Kiro has confirmed via RPC. Both init to `undefined` at spawn; first turn fires the RPC to align Kiro's state with the user's selection even when the user-chosen value matches what spawn-`--agent` set (closes PR #4's first-turn no-op). Behavior: - Single kiro-cli child process per thread. No respawns on agent or model change. - Initial spawn still passes `--agent <name>` as a hint so first turn doesn't need set_mode if user stuck with the default. `--model` is deliberately NOT passed. - Known Kiro-side limitation: switching across model families mid- conversation (e.g., Claude → DeepSeek) can surface Bedrock `ValidationException` from history replay. This is a Kiro product bug to be filed upstream; our client surfaces the error rather than masking it. Tests: - Added regression tests for first-turn set_model/set_mode alignment - Inverted prior PR #4 test whose invariant was "don't respawn on model change" (correct) but failed to assert the set_model RPC - Dropped `--model` spawn-arg tests (invariant reversed) - Agent-respawn test replaced with set_mode-in-session test - Mock agents extended with handleExtRequest("session/set_mode", ...) Verified end-to-end on thread 593e16ae: 1 initialize, 3 set_mode, 4 set_model, 7 prompts, 0 failures, 0 respawns. 935 server tests pass, bun typecheck clean.
1 parent da71e1e commit f0fa43b

4 files changed

Lines changed: 209 additions & 68 deletions

File tree

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

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import { appendFileSync } from "node:fs";
33

44
import * as Effect from "effect/Effect";
5+
import { Schema } from "effect";
56

67
import * as NodeServices from "@effect/platform-node/NodeServices";
78
import * as NodeRuntime from "@effect/platform-node/NodeRuntime";
@@ -248,15 +249,6 @@ const program = Effect.gen(function* () {
248249
),
249250
);
250251

251-
yield* agent.handleSetSessionModel((request) =>
252-
Effect.gen(function* () {
253-
if (typeof request.modelId === "string") {
254-
currentModelId = request.modelId;
255-
}
256-
return {};
257-
}),
258-
);
259-
260252
yield* agent.handleSetSessionConfigOption((request) =>
261253
Effect.gen(function* () {
262254
if (exitOnSetConfigOption) {
@@ -300,6 +292,24 @@ const program = Effect.gen(function* () {
300292
}),
301293
);
302294

295+
yield* agent.handleSetSessionModel((request) =>
296+
Effect.gen(function* () {
297+
if (typeof request.modelId === "string") {
298+
currentModelId = request.modelId;
299+
}
300+
return {};
301+
}),
302+
);
303+
304+
yield* agent.handleExtRequest("session/set_mode", Schema.Unknown, (request: any) =>
305+
Effect.gen(function* () {
306+
if (typeof request.modeId === "string") {
307+
currentModeId = request.modeId;
308+
}
309+
return {};
310+
}),
311+
);
312+
303313
yield* agent.handlePrompt((request) =>
304314
Effect.gen(function* () {
305315
const requestedSessionId = String(request.sessionId ?? sessionId);

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
* Emits: session/update (agent_message_chunk), _kiro.dev/metadata extension.
88
*/
99
import * as Effect from "effect/Effect";
10+
import { Schema } from "effect";
1011

1112
import * as NodeServices from "@effect/platform-node/NodeServices";
1213
import * as NodeRuntime from "@effect/platform-node/NodeRuntime";
@@ -18,6 +19,7 @@ const sessionId = "kiro-mock-session-1";
1819
/** Available models — returned in session responses. */
1920
const availableModels = ["auto", "claude-sonnet-4-20250514", "claude-opus-4-20250918"];
2021
let currentModel = "auto";
22+
let currentMode = "kiro_default";
2123

2224
/** Build the model config option structure used in ACP responses. */
2325
function makeModelConfigOption(model: string) {
@@ -50,6 +52,10 @@ const program = Effect.gen(function* () {
5052
yield* agent.handleCreateSession(() =>
5153
Effect.succeed({
5254
sessionId,
55+
modes: {
56+
currentModeId: currentMode,
57+
availableModes: [],
58+
},
5359
configOptions: [makeModelConfigOption(currentModel)],
5460
}),
5561
);
@@ -67,6 +73,15 @@ const program = Effect.gen(function* () {
6773
}),
6874
);
6975

76+
yield* agent.handleExtRequest("session/set_mode", Schema.Unknown, (request: any) =>
77+
Effect.gen(function* () {
78+
if (typeof request.modeId === "string") {
79+
currentMode = request.modeId;
80+
}
81+
return {};
82+
}),
83+
);
84+
7085
yield* agent.handleSetSessionConfigOption((request) =>
7186
Effect.gen(function* () {
7287
if (request.configId === "model" && "value" in request && typeof request.value === "string") {

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

Lines changed: 107 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -258,11 +258,11 @@ describe("KiroAdapterLive integration", () => {
258258
);
259259

260260
it.effect(
261-
"respawns session with new --agent when sendTurn changes agent mid-session",
261+
"switches agent in-session via session/set_mode without respawning kiro-cli",
262262
() =>
263263
Effect.gen(function* () {
264264
const adapter = yield* KiroAdapter;
265-
const threadId = ThreadId.make("kiro-int-agent-change-1");
265+
const threadId = ThreadId.make("kiro-int-agent-set-1");
266266

267267
yield* adapter.startSession({
268268
threadId,
@@ -281,11 +281,10 @@ describe("KiroAdapterLive integration", () => {
281281
expect(startupArgs).toContain("--agent");
282282
expect(startupArgs).toContain("agent-one");
283283

284-
// Changing the agent mid-session must trigger a respawn because
285-
// --agent is a kiro-cli spawn flag, not an in-session protocol field.
284+
// Agent change MUST NOT respawn the child process — set_mode RPC only.
286285
yield* adapter.sendTurn({
287286
threadId,
288-
input: "hello",
287+
input: "switch agent",
289288
attachments: [],
290289
modelSelection: {
291290
provider: "kiro",
@@ -294,22 +293,19 @@ describe("KiroAdapterLive integration", () => {
294293
},
295294
});
296295

297-
expect(capturedArgs.length).toBeGreaterThan(startupSpawnCount);
298-
const latestArgs = capturedArgs[capturedArgs.length - 1];
299-
expect(latestArgs).toContain("--agent");
300-
expect(latestArgs).toContain("agent-two");
301-
expect(latestArgs).not.toContain("agent-one");
296+
// Critical assertion: spawn count unchanged.
297+
expect(capturedArgs.length).toBe(startupSpawnCount);
302298

303299
yield* adapter.stopSession(threadId);
304300
}).pipe(Effect.scoped, Effect.provide(adapterLayer)),
305301
);
306302

307303
it.effect(
308-
"does not respawn when sendTurn keeps the same agent",
304+
"fires session/set_mode on the first turn when a concrete agent is selected",
309305
() =>
310306
Effect.gen(function* () {
311307
const adapter = yield* KiroAdapter;
312-
const threadId = ThreadId.make("kiro-int-agent-same-1");
308+
const threadId = ThreadId.make("kiro-int-agent-first-turn");
313309

314310
yield* adapter.startSession({
315311
threadId,
@@ -319,20 +315,22 @@ describe("KiroAdapterLive integration", () => {
319315
modelSelection: {
320316
provider: "kiro",
321317
model: "auto",
322-
options: [{ id: "agent", value: "agent-stable" }],
318+
options: [{ id: "agent", value: "ncs-agent" }],
323319
},
324320
});
325321

326322
const spawnCountAfterStart = capturedArgs.length;
327323

324+
// First turn with the same agent that was selected at startSession.
325+
// set_mode should fire to align Kiro's state, without respawn.
328326
yield* adapter.sendTurn({
329327
threadId,
330-
input: "hello",
328+
input: "hi",
331329
attachments: [],
332330
modelSelection: {
333331
provider: "kiro",
334332
model: "auto",
335-
options: [{ id: "agent", value: "agent-stable" }],
333+
options: [{ id: "agent", value: "ncs-agent" }],
336334
},
337335
});
338336

@@ -343,11 +341,51 @@ describe("KiroAdapterLive integration", () => {
343341
);
344342

345343
it.effect(
346-
"switches model in-session without restarting the process",
344+
"fires session/set_model on the first turn when a concrete model is selected",
345+
() =>
346+
Effect.gen(function* () {
347+
const adapter = yield* KiroAdapter;
348+
const threadId = ThreadId.make("kiro-int-model-first-turn");
349+
350+
yield* adapter.startSession({
351+
threadId,
352+
provider: "kiro",
353+
cwd: process.cwd(),
354+
runtimeMode: "full-access",
355+
modelSelection: {
356+
provider: "kiro",
357+
model: "claude-opus-4.6",
358+
},
359+
});
360+
361+
// First turn with the same model that was selected at startSession.
362+
// Without the activeModel-based gate, set_model would not fire here
363+
// and Kiro would stay on its internal default (often a deprecated
364+
// preview like claude-opus-4.6-1m).
365+
yield* adapter.sendTurn({
366+
threadId,
367+
input: "first turn",
368+
attachments: [],
369+
modelSelection: {
370+
provider: "kiro",
371+
model: "claude-opus-4.6",
372+
},
373+
});
374+
375+
const sessions = yield* adapter.listSessions();
376+
const currentSession = sessions.find((s) => s.threadId === threadId);
377+
expect(currentSession?.model).toBe("claude-opus-4.6");
378+
379+
yield* adapter.stopSession(threadId);
380+
}).pipe(Effect.scoped, Effect.provide(adapterLayer)),
381+
);
382+
383+
it.effect(
384+
"switches model in-session via session/set_model without respawning kiro-cli",
347385
() =>
348386
Effect.gen(function* () {
349387
const adapter = yield* KiroAdapter;
350-
const threadId = ThreadId.make("kiro-int-model-switch-1");
388+
const threadId = ThreadId.make("kiro-int-model-set-1");
351389

352390
yield* adapter.startSession({
353391
threadId,
@@ -359,31 +397,76 @@ describe("KiroAdapterLive integration", () => {
359397

360398
const spawnCountAfterStart = capturedArgs.length;
361399

362-
// First turn with default model
400+
// First turn on model "auto" — no set_model call expected
363401
yield* adapter.sendTurn({
364402
threadId,
365403
input: "first turn",
366404
attachments: [],
367405
});
368406

369-
// Second turn with different model — should NOT respawn
407+
// Second turn switches model — should call set_model, NOT respawn
370408
yield* adapter.sendTurn({
371409
threadId,
372-
input: "second turn after model switch",
410+
input: "second turn with switched model",
373411
attachments: [],
374412
modelSelection: {
375413
provider: "kiro",
376-
model: "claude-sonnet-4-20250514",
414+
model: "claude-sonnet-4.6",
415+
},
416+
});
417+
418+
// Critical assertion: no new spawn fired. Model change is an in-session
419+
// RPC, not a process restart.
420+
expect(capturedArgs.length).toBe(spawnCountAfterStart);
421+
422+
// Session state reflects the new model.
423+
const sessions = yield* adapter.listSessions();
424+
const currentSession = sessions.find((s) => s.threadId === threadId);
425+
expect(currentSession?.model).toBe("claude-sonnet-4.6");
426+
427+
yield* adapter.stopSession(threadId);
428+
}).pipe(Effect.scoped, Effect.provide(adapterLayer)),
429+
);
430+
431+
it.effect(
432+
"does not respawn when sendTurn keeps the same agent",
433+
() =>
434+
Effect.gen(function* () {
435+
const adapter = yield* KiroAdapter;
436+
const threadId = ThreadId.make("kiro-int-agent-same-1");
437+
438+
yield* adapter.startSession({
439+
threadId,
440+
provider: "kiro",
441+
cwd: process.cwd(),
442+
runtimeMode: "full-access",
443+
modelSelection: {
444+
provider: "kiro",
445+
model: "auto",
446+
options: [{ id: "agent", value: "agent-stable" }],
447+
},
448+
});
449+
450+
const spawnCountAfterStart = capturedArgs.length;
451+
452+
yield* adapter.sendTurn({
453+
threadId,
454+
input: "hello",
455+
attachments: [],
456+
modelSelection: {
457+
provider: "kiro",
458+
model: "auto",
459+
options: [{ id: "agent", value: "agent-stable" }],
377460
},
378461
});
379462

380-
// Session should NOT have been restarted — only one spawn
381463
expect(capturedArgs.length).toBe(spawnCountAfterStart);
382464

383465
yield* adapter.stopSession(threadId);
384466
}).pipe(Effect.scoped, Effect.provide(adapterLayer)),
385467
);
386468

469+
387470
it.effect(
388471
"updates session.model immediately after in-session model switch",
389472
() =>
@@ -401,7 +484,7 @@ describe("KiroAdapterLive integration", () => {
401484

402485
expect(session.model).toBe("auto");
403486

404-
// Collect turn.started events to verify model is correct
487+
// Collect turn.started events to verify model is correct after in-session switch
405488
const eventsFiber = yield* adapter.streamEvents.pipe(
406489
Stream.filter(
407490
(event) => event.type === "turn.started" || event.type === "turn.completed",
@@ -411,7 +494,7 @@ describe("KiroAdapterLive integration", () => {
411494
Effect.forkChild,
412495
);
413496

414-
// Send turn with new model
497+
// Send turn with new model (triggers in-session set_model RPC)
415498
const turn = yield* adapter.sendTurn({
416499
threadId,
417500
input: "switch model turn",

0 commit comments

Comments
 (0)