package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/pelletier/go-toml/v2"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/provider"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

const kimiTestSessionID = "baf8f065-a3b1-45f4-9a3d-b760b3db80d1"

func TestKimiPrepareLaunchUsesWireSessionAndWorkDir(t *testing.T) {
	workDir := t.TempDir()
	prepared, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir: workDir,
		Env:     map[string]string{"KIMI_SHARE_DIR": filepath.Join(workDir, "share")},
	}, kimiTestSessionID)
	require.NoError(t, err)
	require.Equal(t, []string{"--wire", "--yolo", "--session", kimiTestSessionID, "--work-dir", workDir}, prepared.Args)
	require.Equal(t, filepath.Join(workDir, "share"), prepared.ShareDir)
	require.NotContains(t, strings.Join(prepared.Env, "\n"), "KIMI_CODE_HOME=")
	require.Empty(t, prepared.AgentFile)
	require.Empty(t, prepared.MCPConfigFile)
}

func TestKimiPrepareLaunchToolRestrictedFailsClosedWithoutLosingOAuth(t *testing.T) {
	workDir := t.TempDir()
	shareDir := filepath.Join(t.TempDir(), "user-oauth-home")
	require.NoError(t, os.MkdirAll(filepath.Join(shareDir, "credentials"), 0o700))
	require.NoError(t, os.WriteFile(filepath.Join(shareDir, "credentials", "oauth.json"), []byte(`{"access_token":"test"}`), 0o600))
	require.NoError(t, os.WriteFile(filepath.Join(shareDir, "credentials", "unrelated.json"), []byte(`{"secret":"must-not-copy"}`), 0o600))
	require.NoError(t, os.WriteFile(filepath.Join(shareDir, "config.toml"), []byte(`
default_model = "kimi"
default_plan_mode = true
default_editor = "sh -c unsafe"
extra_skill_dirs = ["/tmp/ambient-skills"]

[providers.selected]
type = "kimi"
base_url = "https://selected.example"
api_key = "selected-secret"

[models.kimi]
provider = "selected"
model = "kimi"
max_context_size = 262144

[providers.unrelated]
type = "anthropic"
base_url = "https://unrelated.example"
api_key = "must-not-copy"

[models.unrelated]
provider = "unrelated"
model = "claude"
max_context_size = 262144

[services.moonshot_search]
base_url = "https://search.example"
api_key = "must-not-copy"

[[hooks]]
event = "SessionStart"
command = "touch /tmp/unsafe"
`), 0o600))
	require.NoError(t, os.MkdirAll(filepath.Join(shareDir, "plugins", "ambient"), 0o700))
	require.NoError(t, os.WriteFile(filepath.Join(shareDir, "plugins", "ambient", "plugin.json"), []byte(`{}`), 0o600))
	t.Setenv("KIMI_SHARE_DIR", shareDir)
	prepared, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir:            workDir,
		DisableNativeTools: true,
		PermissionMode:     types.AgentPermissionModeFullAccess,
	}, kimiTestSessionID)
	require.NoError(t, err)

	agentYAML, err := os.ReadFile(prepared.AgentFile)
	require.NoError(t, err)
	require.Contains(t, string(agentYAML), "tools: []")
	require.Contains(t, string(agentYAML), "subagents: {}")
	mcpJSON, err := os.ReadFile(prepared.MCPConfigFile)
	require.NoError(t, err)
	require.JSONEq(t, `{"mcpServers":{}}`, string(mcpJSON))
	require.Contains(t, prepared.Args, "--agent-file")
	require.Contains(t, prepared.Args, "--mcp-config-file")
	require.Contains(t, prepared.Args, "--skills-dir")
	skillsArg := prepared.Args[argumentIndex(t, prepared.Args, "--skills-dir")+1]
	require.Equal(t, filepath.Join(prepared.ShareDir, "agentteam-skills-empty"), skillsArg)
	require.DirExists(t, skillsArg)
	require.Contains(t, strings.Join(prepared.Env, "\n"), "KIMI_SHARE_DIR="+prepared.ShareDir)
	require.NotContains(t, strings.Join(prepared.Env, "\n"), "KIMI_CODE_HOME=")
	require.True(t, strings.HasPrefix(prepared.AgentFile, filepath.Join(prepared.ShareDir, "agentteam-launch")+string(filepath.Separator)))
	require.Contains(t, prepared.Args, "--yolo")
	require.Equal(t, filepath.Join(filepath.Dir(workDir), "kimi-share"), prepared.ShareDir)
	require.FileExists(t, filepath.Join(prepared.ShareDir, "credentials", "oauth.json"))
	require.NoFileExists(t, filepath.Join(prepared.ShareDir, "credentials", "unrelated.json"))
	require.NoDirExists(t, filepath.Join(prepared.ShareDir, "plugins"))
	isolatedConfig, err := os.ReadFile(filepath.Join(prepared.ShareDir, "config.toml"))
	require.NoError(t, err)
	var decodedConfig map[string]any
	require.NoError(t, toml.Unmarshal(isolatedConfig, &decodedConfig))
	require.Equal(t, "kimi", decodedConfig["default_model"])
	require.Contains(t, string(isolatedConfig), `default_model = 'kimi'`)
	require.NotContains(t, string(isolatedConfig), "hooks")
	require.NotContains(t, string(isolatedConfig), "extra_skill_dirs")
	require.NotContains(t, string(isolatedConfig), "default_editor")
	require.NotContains(t, string(isolatedConfig), "default_plan_mode")
	require.NotContains(t, string(isolatedConfig), "moonshot_search")
	require.NotContains(t, string(isolatedConfig), "must-not-copy")
	require.Contains(t, string(isolatedConfig), "selected-secret")
	require.NotContains(t, string(isolatedConfig), "unrelated.example")
	prepared.Cleanup()
	require.DirExists(t, prepared.ShareDir)
	require.NoDirExists(t, filepath.Join(prepared.ShareDir, "credentials"))
	require.NoFileExists(t, filepath.Join(prepared.ShareDir, "config.toml"))
	require.NoDirExists(t, filepath.Dir(prepared.AgentFile))
}

func TestKimiPrepareLaunchCleanupComposesIsolatedShareAndLaunchConfig(t *testing.T) {
	workDir := t.TempDir()
	prepared, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir:            workDir,
		DisableNativeTools: true,
		PermissionMode:     types.AgentPermissionModeFullAccess,
		MCPs: []types.NamedMCPConfig{
			{Name: "knowledge", URL: "https://agentteam.example/mcp"},
		},
	}, kimiTestSessionID)
	require.NoError(t, err)
	launchConfigDir := filepath.Dir(prepared.AgentFile)
	require.DirExists(t, prepared.ShareDir)
	require.DirExists(t, launchConfigDir)
	sessionDir := kimiTestSessionDir(t, prepared.ShareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))

	prepared.Cleanup()

	require.DirExists(t, prepared.ShareDir)
	require.DirExists(t, sessionDir)
	require.NoDirExists(t, launchConfigDir)
}

func TestKimiPrepareLaunchProviderFailureCleansIsolatedCredentialsAndLaunchDir(t *testing.T) {
	workDir := t.TempDir()
	sourceShareDir := filepath.Join(t.TempDir(), "source-share")
	require.NoError(t, os.MkdirAll(filepath.Join(sourceShareDir, "credentials"), 0o700))
	require.NoError(t, os.WriteFile(filepath.Join(sourceShareDir, "credentials", "oauth.json"), []byte(`{"access_token":"test"}`), 0o600))
	require.NoError(t, os.WriteFile(filepath.Join(sourceShareDir, "config.toml"), []byte(`
default_model = "kimi"
[providers.kimi]
type = "kimi"
base_url = "https://example.com"
api_key = "secret"
[models.kimi]
provider = "kimi"
model = "kimi"
max_context_size = 262144
	`), 0o600))
	isolatedShareDir := filepath.Join(filepath.Dir(workDir), "kimi-share")
	var providerConfigPath string
	originalWrite := writeKimiProviderConfigFile
	writeKimiProviderConfigFile = func(path string, data []byte, mode os.FileMode) error {
		providerConfigPath = path
		require.NoError(t, os.WriteFile(path, data, mode))
		return os.ErrPermission
	}
	t.Cleanup(func() { writeKimiProviderConfigFile = originalWrite })

	_, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir:            workDir,
		Env:                map[string]string{"KIMI_SHARE_DIR": sourceShareDir},
		DisableNativeTools: true,
		PermissionMode:     types.AgentPermissionModeFullAccess,
		Provider: &drivercontract.ProviderConfig{
			APIKey:   "sk-test",
			BaseURL:  "https://llm.example.com/v1",
			WireAPIs: []string{provider.WireAPIResponses},
		},
		Model: "kimi",
	}, kimiTestSessionID)
	require.Error(t, err)
	require.NotEmpty(t, providerConfigPath)
	require.NoFileExists(t, filepath.Join(isolatedShareDir, "config.toml"))
	require.NoDirExists(t, filepath.Join(isolatedShareDir, "credentials"))
	require.NoDirExists(t, filepath.Dir(providerConfigPath))
}

func TestKimiRejectsRemovedKnowledgeQAPermissionMode(t *testing.T) {
	_, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir:        t.TempDir(),
		PermissionMode: "knowledge_qa_scratch",
	}, kimiTestSessionID)
	require.ErrorContains(t, err, `unsupported agent permission mode "knowledge_qa_scratch"`)
}

func TestKimiToolRestrictedResumeUsesStableIsolatedSessionWithoutCredentialResidue(t *testing.T) {
	workDir := t.TempDir()
	sourceShareDir := filepath.Join(t.TempDir(), "source-share")
	require.NoError(t, os.MkdirAll(filepath.Join(sourceShareDir, "credentials"), 0o700))
	require.NoError(t, os.WriteFile(filepath.Join(sourceShareDir, "credentials", "oauth.json"), []byte(`{}`), 0o600))
	cfg := drivercontract.LaunchConfig{
		WorkDir:            workDir,
		Env:                map[string]string{"KIMI_SHARE_DIR": sourceShareDir},
		DisableNativeTools: true,
		PermissionMode:     types.AgentPermissionModeFullAccess,
		SessionID:          kimiTestSessionID,
	}
	prepared, err := prepareKimiLaunch(cfg, kimiTestSessionID)
	require.NoError(t, err)
	sessionDir := kimiTestSessionDir(t, prepared.ShareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	prepared.Cleanup()
	require.NoDirExists(t, filepath.Join(prepared.ShareDir, "credentials"))

	sessionID, resumed, err := resolveKimiLaunchSession(cfg)
	require.NoError(t, err)
	require.True(t, resumed)
	require.Equal(t, kimiTestSessionID, sessionID)
}

func TestKimiLaunchSessionFallsBackToFreshSession(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	cfg := testLaunchConfig(drivercontract.LaunchConfig{WorkDir: workDir, Env: map[string]string{"KIMI_SHARE_DIR": shareDir}, SessionID: kimiTestSessionID}, testLaunchSurfaces{KimiInitialPrompt: "fresh", FreshPrompt: "fresh", ResumePrompt: "resume"})
	sessionID, resumed, err := resolveKimiLaunchSession(cfg)
	require.NoError(t, err)
	require.False(t, resumed)
	require.NotEqual(t, kimiTestSessionID, sessionID)
	require.Contains(t, kimiInitialPrompt(cfg, resumed), "fresh")
	require.NotContains(t, kimiInitialPrompt(cfg, resumed), "resume")

}

func TestKimiLaunchSessionIncompleteTurnFallsBack(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	writeKimiLines(t, filepath.Join(sessionDir, "wire.jsonl"), []string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "unfinished"}),
	})
	cfg := drivercontract.LaunchConfig{
		WorkDir:   workDir,
		Env:       map[string]string{"KIMI_SHARE_DIR": shareDir},
		SessionID: kimiTestSessionID,
	}
	sessionID, resumed, err := resolveKimiLaunchSession(cfg)
	require.NoError(t, err)
	require.False(t, resumed)
	require.NotEqual(t, kimiTestSessionID, sessionID)

}

func TestKimiLaunchSessionRejectsUnsafeExistingSessionPath(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(filepath.Dir(sessionDir), 0o700))
	require.NoError(t, os.WriteFile(sessionDir, []byte("not a directory"), 0o600))

	_, _, err := resolveKimiLaunchSession(drivercontract.LaunchConfig{
		WorkDir:   workDir,
		Env:       map[string]string{"KIMI_SHARE_DIR": shareDir},
		SessionID: kimiTestSessionID,
	})
	require.Error(t, err)
}

func TestKimiPrepareLaunchUsesShellShareDirForSessionPathParity(t *testing.T) {
	workDir := t.TempDir()
	shareDir := filepath.Join(t.TempDir(), "shell-share")
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{"KIMI_SHARE_DIR": shareDir})
	t.Cleanup(func() { runtimeconfig.SetShellEnvForRuntimeResolution(nil) })
	prepared, err := prepareKimiLaunch(drivercontract.LaunchConfig{WorkDir: workDir}, kimiTestSessionID)
	require.NoError(t, err)
	require.Contains(t, strings.Join(prepared.Env, "\n"), "KIMI_SHARE_DIR="+shareDir)
	sessionDir, err := kimiSessionDir(drivercontract.LaunchConfig{WorkDir: workDir}, kimiTestSessionID, false)
	require.NoError(t, err)
	require.True(t, strings.HasPrefix(sessionDir, filepath.Join(shareDir, "sessions")+string(filepath.Separator)))
}

func TestKimiPrepareLaunchWritesConfiguredMCPs(t *testing.T) {
	workDir := t.TempDir()
	prepared, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir: workDir,
		MCPs: []types.NamedMCPConfig{
			{Name: "agentteam", URL: "https://agentteam.example/mcp", Headers: map[string]string{"Authorization": "Bearer test"}},
			{Name: "local", Transport: types.MCPTransportStdio, Command: "agentteam-mcp", Args: []string{"serve"}, Env: map[string]string{"A": "B"}, CWD: workDir},
		},
	}, kimiTestSessionID)
	require.NoError(t, err)
	require.Empty(t, prepared.AgentFile)
	require.Contains(t, prepared.Args, "--mcp-config-file")
	payload, err := os.ReadFile(prepared.MCPConfigFile)
	require.NoError(t, err)
	require.JSONEq(t, `{
		"mcpServers": {
			"agentteam": {"url":"https://agentteam.example/mcp","headers":{"Authorization":"Bearer test"}},
			"local": {"command":"agentteam-mcp","args":["serve"],"env":{"A":"B"},"cwd":"`+workDir+`"}
		}
	}`, string(payload))
}

func TestKimiInitialPromptIncludesSplitInstructionsAndResumeInput(t *testing.T) {
	fresh := kimiInitialPrompt(testLaunchConfig(drivercontract.LaunchConfig{}, testLaunchSurfaces{KimiInitialPrompt: "base contract\n\ndeveloper contract\n\nfresh turn", FreshPrompt: "fresh turn", BaseInstructions: "base contract", DeveloperInstructions: "developer contract"}), false)
	require.Contains(t, fresh, "base contract")
	require.Contains(t, fresh, "developer contract")
	require.Contains(t, fresh, "fresh turn")

	resume := kimiInitialPrompt(testLaunchConfig(drivercontract.LaunchConfig{SessionID: kimiTestSessionID}, testLaunchSurfaces{KimiInitialPrompt: "base contract\n\nresume turn", FreshPrompt: "fresh turn", ResumePrompt: "resume turn", BaseInstructions: "base contract"}), true)
	require.Contains(t, resume, "base contract")
	require.Contains(t, resume, "resume turn")
	require.NotContains(t, resume, "fresh turn")
}

func TestKimiPrepareLaunchWritesDocumentedProviderConfig(t *testing.T) {
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{"AGENTTEAM_TEST_KIMI_KEY": "sk-shell"})
	t.Cleanup(func() { runtimeconfig.SetShellEnvForRuntimeResolution(nil) })
	prepared, err := prepareKimiLaunch(drivercontract.LaunchConfig{
		WorkDir: t.TempDir(),
		Env:     map[string]string{"EXISTING": "1", "KIMI_SHARE_DIR": t.TempDir()},
		Model:   "agentteam/gpt-5.5",
		Provider: &drivercontract.ProviderConfig{
			ID: "agentteam", Name: "AgentTeam Gateway", BaseURL: "https://llm.example.com/v1",
			EnvKey: "AGENTTEAM_TEST_KIMI_KEY", Headers: map[string]string{"X-AgentTeam": "1"},
			WireAPIs: []string{provider.WireAPIResponses},
		},
	}, kimiTestSessionID)
	require.NoError(t, err)
	require.Contains(t, prepared.Args, "--config-file")
	require.Contains(t, prepared.Args, "--model")
	require.Equal(t, "agentteam-runtime", prepared.Args[argumentIndex(t, prepared.Args, "--model")+1])
	configPath := prepared.Args[argumentIndex(t, prepared.Args, "--config-file")+1]
	configData, err := os.ReadFile(configPath)
	require.NoError(t, err)
	var config map[string]any
	require.NoError(t, toml.Unmarshal(configData, &config))
	require.Equal(t, "agentteam-runtime", config["default_model"])
	require.Contains(t, string(configData), `type = 'openai_responses'`)
	require.Contains(t, string(configData), `base_url = 'https://llm.example.com/v1'`)
	require.Contains(t, string(configData), `api_key = 'sk-shell'`)
	require.Contains(t, string(configData), `model = 'gpt-5.5'`)
	require.Contains(t, string(configData), `X-AgentTeam = '1'`)
	require.NotContains(t, strings.Join(prepared.Env, "\n"), "KIMI_MODEL_")
}

func TestKimiWireServerRequestsFailClosed(t *testing.T) {
	var stdin bytes.Buffer
	d := &KimiWireDriver{stdin: nopWriteCloser{Writer: &stdin}, pending: map[string]chan kimiWireEnvelope{}}
	d.handleServerRequest(kimiWireEnvelope{
		ID: json.RawMessage(`"approval-1"`), Method: "request",
		Params: json.RawMessage(`{"type":"ApprovalRequest","payload":{"id":"approval-inner"}}`),
	})
	d.handleServerRequest(kimiWireEnvelope{
		ID: json.RawMessage(`"question-1"`), Method: "request",
		Params: json.RawMessage(`{"type":"QuestionRequest","payload":{"id":"question-inner"}}`),
	})
	d.handleServerRequest(kimiWireEnvelope{
		ID: json.RawMessage(`"tool-1"`), Method: "request",
		Params: json.RawMessage(`{"type":"ToolCallRequest","payload":{"id":"tool-inner"}}`),
	})
	lines := strings.Split(strings.TrimSpace(stdin.String()), "\n")
	require.Len(t, lines, 3)
	require.JSONEq(t, `{"jsonrpc":"2.0","id":"approval-1","result":{"request_id":"approval-inner","response":"reject","feedback":"AgentTeam does not grant interactive Kimi Wire approvals"}}`, lines[0])
	require.JSONEq(t, `{"jsonrpc":"2.0","id":"question-1","result":{"request_id":"question-inner","answers":{}}}`, lines[1])
	require.JSONEq(t, `{"jsonrpc":"2.0","id":"tool-1","error":{"code":-32601,"message":"unsupported Kimi Wire request type \"ToolCallRequest\""}}`, lines[2])
}

func TestKimiInitializeResultRequiresExactWireProtocol(t *testing.T) {
	require.NoError(t, validateKimiInitializeResult(json.RawMessage(`{
		"protocol_version":"1.10",
		"server":{"name":"Kimi Code CLI","version":"1.49.0"}
	}`)))
	require.Error(t, validateKimiInitializeResult(json.RawMessage(`{
		"protocol_version":"1.9",
		"server":{"name":"Kimi Code CLI","version":"1.49.0"}
	}`)))
	require.Error(t, validateKimiInitializeResult(json.RawMessage(`{"protocol_version":"1.10"}`)))
}

func TestKimiMutatingRequestTransportCloseIsAmbiguous(t *testing.T) {
	var stdin bytes.Buffer
	driver := newKimiTestDriver(1)
	driver.stdin = nopWriteCloser{Writer: &stdin}
	done := make(chan error, 1)
	go func() {
		_, err := driver.request(context.Background(), "prompt", map[string]any{"user_input": "hello"})
		done <- err
	}()
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.pending) == 1
	}, time.Second, time.Millisecond)
	driver.failPending(io.ErrUnexpectedEOF)
	require.ErrorIs(t, <-done, errKimiWireTransportAmbiguous)
}

func TestKimiReadOnlyRequestTransportCloseRemainsRetryable(t *testing.T) {
	var stdin bytes.Buffer
	driver := newKimiTestDriver(1)
	driver.stdin = nopWriteCloser{Writer: &stdin}
	done := make(chan error, 1)
	go func() {
		_, err := driver.request(context.Background(), "initialize", map[string]any{})
		done <- err
	}()
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.pending) == 1
	}, time.Second, time.Millisecond)
	driver.failPending(io.ErrUnexpectedEOF)
	require.NotErrorIs(t, <-done, errKimiWireTransportAmbiguous)
}

func TestKimiTranslatedOutputCloseReleasesBlockedEmitter(t *testing.T) {
	output := newKimiTranslatedOutput(1)
	require.True(t, output.emit("first"))
	emitDone := make(chan bool, 1)
	go func() {
		emitDone <- output.emit("blocked")
	}()
	select {
	case <-emitDone:
		t.Fatal("emitter did not back-pressure on a full queue")
	case <-time.After(20 * time.Millisecond):
	}

	closeDone := make(chan struct{})
	go func() {
		output.close()
		close(closeDone)
	}()
	select {
	case emitted := <-emitDone:
		require.False(t, emitted)
	case <-time.After(time.Second):
		t.Fatal("closing translated output did not release blocked emitter")
	}
	select {
	case <-closeDone:
	case <-time.After(time.Second):
		t.Fatal("translated output close waited forever for emitter")
	}
	_, ok := <-output.lines
	require.True(t, ok, "the admitted line remains buffered until drained")
	_, ok = <-output.lines
	require.False(t, ok)
}

func TestKimiTranslatedOutputCloseUnblocksMissingPipeConsumer(t *testing.T) {
	output := newKimiTranslatedOutput(1)
	reader, writer := io.Pipe()
	defer reader.Close()
	output.attachWriter(writer)
	flusherDone := make(chan struct{})
	go func() {
		defer close(flusherDone)
		defer writer.Close()
		for line := range output.lines {
			_, _ = io.WriteString(writer, line+"\n")
		}
	}()
	require.True(t, output.emit("blocked on pipe"))

	require.Eventually(t, func() bool {
		output.mu.Lock()
		defer output.mu.Unlock()
		return len(output.lines) == 0
	}, time.Second, time.Millisecond)
	output.close()
	select {
	case <-flusherDone:
	case <-time.After(time.Second):
		t.Fatal("translated output close left flusher blocked on an unread pipe")
	}
}

func TestKimiPromptPublishesTerminalOnlyAfterRPCOutcome(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.promptInFlight = true
	driver.handleNotification("event", json.RawMessage(`{"type":"TurnEnd","payload":{}}`))
	select {
	case line := <-driver.lines:
		t.Fatalf("TurnEnd published before prompt response: %s", line)
	default:
	}
	driver.finishPrompt(kimiPromptOutcome{ready: true, errMessage: "provider failed", stopReason: "error"})
	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, consts.ParsedEventKindError, events[0].Kind)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[1].Kind)
	require.Equal(t, "error", events[1].StopReason)
}

func TestKimiPromptResponseBeforeTurnEndStillPublishesSingleTerminal(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.promptInFlight = true
	driver.finishPrompt(kimiPromptOutcome{ready: true, stopReason: consts.ParsedEventStopReasonEndTurn})
	select {
	case line := <-driver.lines:
		t.Fatalf("prompt response published before TurnEnd: %s", line)
	default:
	}
	driver.handleNotification("event", json.RawMessage(`{"type":"TurnEnd","payload":{}}`))
	events := drainKimiParsedEvents(t, driver, 1)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[0].Kind)
	require.Equal(t, consts.ParsedEventStopReasonEndTurn, events[0].StopReason)
	select {
	case line := <-driver.lines:
		t.Fatalf("terminal published twice: %s", line)
	default:
	}
}

func TestKimiPromptPublishesTerminalWhenTransportClosesAfterRPCFailure(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.promptInFlight = true
	driver.closeTransport()
	select {
	case line := <-driver.lines:
		t.Fatalf("transport close published before prompt outcome: %s", line)
	default:
	}
	driver.finishPrompt(kimiPromptOutcome{ready: true, errMessage: "transport closed", stopReason: "error"})
	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, consts.ParsedEventKindError, events[0].Kind)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[1].Kind)
	require.Equal(t, "error", events[1].StopReason)
}

func TestKimiTransportCloseFlushesReadyPromptOutcomeOnce(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.promptOutcome = kimiPromptOutcome{ready: true, errMessage: "transport closed", stopReason: "error"}
	driver.promptInFlight = false
	driver.turnActive = true
	driver.closeTransport()
	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, consts.ParsedEventKindError, events[0].Kind)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[1].Kind)
	driver.closeTransport()
	select {
	case line := <-driver.lines:
		t.Fatalf("terminal published twice: %s", line)
	default:
	}
}

func TestKimiReadLoopMarksTransportClosedBeforeWakingPrompt(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.stdin = nopWriteCloser{}
	driver.runtimeCtx = context.Background()
	delivered := make(chan error, 1)
	go func() {
		delivered <- driver.DeliverRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "transport-close"}, "hello"), kimiTestSessionID)
	}()
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return driver.promptInFlight && len(driver.pending) == 1
	}, time.Second, time.Millisecond)
	driver.readLoop(strings.NewReader(""), func() { driver.closeTransport() })
	require.ErrorIs(t, <-delivered, errKimiWireTransportAmbiguous)
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return driver.promptOutcome.ready
	}, time.Second, time.Millisecond)
	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, consts.ParsedEventKindError, events[0].Kind)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[1].Kind)
}

func TestKimiPromptFailureBeforeTurnBeginWakesAdmissionAmbiguous(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.runtimeCtx = context.Background()
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		require.Equal(t, "prompt", method)
		return nil, errors.New("prompt rejected")
	}
	err := driver.DeliverRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "prompt-rejected"}, "hello"), kimiTestSessionID)
	require.ErrorIs(t, err, errKimiWireTransportAmbiguous)
	require.ErrorContains(t, err, "prompt rejected")
}

func TestKimiPromptAdmissionCancellationIsAmbiguous(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.runtimeCtx = context.Background()
	release := make(chan struct{})
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		require.Equal(t, "prompt", method)
		<-release
		return json.RawMessage(`{"status":"finished"}`), nil
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	err := driver.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "prompt-cancelled"}, "hello"), kimiTestSessionID)
	close(release)
	require.ErrorIs(t, err, errKimiWireTransportAmbiguous)
	require.ErrorIs(t, err, context.Canceled)
}

func TestKimiSteerTransportCloseBeforeAppliedIsAmbiguous(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		require.Equal(t, "steer", method)
		return json.RawMessage(`{"status":"steered"}`), nil
	}
	done := make(chan error, 1)
	go func() {
		done <- driver.DeliverAckedBusyRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "steer-close"}, "hello"), kimiTestSessionID)
	}()
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return driver.steerWaiters["steer-close"] != nil
	}, time.Second, time.Millisecond)
	driver.closeTransport()
	require.ErrorIs(t, <-done, errKimiWireTransportAmbiguous)
}

func TestKimiWireEventsMapLifecycleContentToolPlanAndUsage(t *testing.T) {
	d := newKimiTestDriver(32)
	d.promptInFlight = true
	d.handleNotification("event", json.RawMessage(`{"type":"TurnBegin","payload":{"user_input":"hello\n\n<!-- agentteam-runtime-input:idle-input -->"}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"ContentPart","payload":{"type":"think","think":"reason"}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"ContentPart","payload":{"type":"text","text":"answer"}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"ToolCall","payload":{"type":"function","id":"tc-1","function":{"name":"Shell","arguments":"{\"command\":\"pwd\"}"}}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"ToolResult","payload":{"tool_call_id":"tc-1","return_value":{"is_error":false,"output":"done","message":"ok","display":[]}}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"PlanDisplay","payload":{"content":"1. inspect\n2. fix","file_path":"/tmp/plan.md"}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"StatusUpdate","payload":{"token_usage":{"input_other":11,"input_cache_read":3,"input_cache_creation":2,"output":7}}}`))
	d.handleNotification("event", json.RawMessage(`{"type":"TurnEnd","payload":{}}`))
	d.finishPrompt(kimiPromptOutcome{ready: true, stopReason: consts.ParsedEventStopReasonEndTurn})

	events := drainKimiParsedEvents(t, d, 9)
	require.Equal(t, consts.ParsedEventKindThinking, events[0].Kind)
	require.Equal(t, true, events[0].Metadata[codexAuthoritativeTurnStartMetadataKey])
	require.Equal(t, consts.ParsedEventKindRuntimeInputAcked, events[1].Kind)
	require.Equal(t, "idle-input", events[1].RuntimeInputID)
	require.Equal(t, "reason", events[2].Text)
	require.Equal(t, "answer", events[3].Text)
	require.Equal(t, consts.ParsedEventKindToolCall, events[4].Kind)
	require.Equal(t, consts.ParsedEventKindToolOutput, events[5].Kind)
	require.Contains(t, events[6].Text, "inspect")
	require.Equal(t, uint64(16), events[7].InputTokens)
	require.Equal(t, uint64(3), events[7].CachedInputTokens)
	require.Equal(t, uint64(7), events[7].OutputTokens)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[8].Kind)
}

func TestKimiWireMapsCompactionMCPAndSubagentActivity(t *testing.T) {
	driver := newKimiTestDriver(8)
	driver.handleNotification("event", json.RawMessage(`{"type":"CompactionBegin","payload":{}}`))
	driver.handleNotification("event", json.RawMessage(`{"type":"MCPLoadingBegin","payload":{}}`))
	driver.handleNotification("event", json.RawMessage(`{"type":"StatusUpdate","payload":{"mcp_status":{"loading":true,"connected":0,"total":1,"tools":0}}}`))
	driver.handleNotification("event", json.RawMessage(`{"type":"SubagentEvent","payload":{"agent_id":"sub-1","subagent_type":"coder","event":{"type":"StepBegin","payload":{"n":1}}}}`))
	events := drainKimiParsedEvents(t, driver, 4)
	require.Equal(t, consts.ParsedEventKindRuntimeCompacting, events[0].Kind)
	require.Equal(t, consts.ParsedEventKindThinking, events[1].Kind)
	require.Equal(t, consts.ParsedEventKindThinking, events[2].Kind)
	require.Equal(t, consts.ParsedEventKindSubagentActivity, events[3].Kind)
	require.Equal(t, "sub-1", events[3].Metadata["agent_id"])
}

func TestKimiWireMapsCompletedCompactionBoundary(t *testing.T) {
	driver := newKimiTestDriver(2)
	driver.handleNotification("event", json.RawMessage(`{"type":"CompactionEnd","payload":{}}`))
	driver.handleNotification("event", json.RawMessage(`{"type":"StatusUpdate","payload":{"message_id":"post-compact","token_usage":{"input_other":10,"input_cache_read":0,"input_cache_creation":0,"output":5}}}`))

	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, consts.ParsedEventKindThinking, events[0].Kind)
	require.True(t, parsedBoolMetadata(events[0].Metadata, contextCompactionConfirmedMetadataKey))
	require.Equal(t, consts.ParsedEventKindUsage, events[1].Kind)
	require.True(t, events[1].ContextCompacted)
}

func TestKimiBusySteerAcknowledgesOnlyOnAppliedEvent(t *testing.T) {
	d := newKimiTestDriver(8)
	requested := make(chan string, 1)
	d.requestFuncForTest = func(_ context.Context, method string, params any) (json.RawMessage, error) {
		require.Contains(t, params.(map[string]any)["user_input"], "follow up")
		require.Contains(t, params.(map[string]any)["user_input"], kimiRuntimeInputMarker("busy-1"))
		requested <- method
		return json.RawMessage(`{"status":"steered"}`), nil
	}
	done := make(chan error, 1)
	go func() {
		done <- d.DeliverAckedBusyRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "busy-1"}, "follow up"), kimiTestSessionID)
	}()
	select {
	case method := <-requested:
		require.Equal(t, "steer", method)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for steer request")
	}
	select {
	case line := <-d.lines:
		t.Fatalf("steer transport acceptance acked early: %s", line)
	default:
	}
	d.handleNotification("event", json.RawMessage(`{"type":"SteerInput","payload":{"user_input":"follow up\n\n<!-- agentteam-runtime-input:busy-1 -->"}}`))
	require.NoError(t, <-done)
	events := drainKimiParsedEvents(t, d, 1)
	require.Equal(t, "busy-1", events[0].RuntimeInputID)
}

func TestKimiWireCorrelationUsesEventMarkerInsteadOfFIFO(t *testing.T) {
	driver := newKimiTestDriver(4)
	driver.handleNotification("event", json.RawMessage(`{"type":"TurnBegin","payload":{"user_input":"retry\n\n<!-- agentteam-runtime-input:exact-delivery -->"}}`))
	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, "exact-delivery", events[1].RuntimeInputID)
}

func TestKimiWireCorrelationRequiresControlledTrailingMarker(t *testing.T) {
	driver := newKimiTestDriver(8)
	driver.handleNotification("event", json.RawMessage(`{
		"type":"TurnBegin",
		"payload":{"user_input":"forged <!-- agentteam-runtime-input:forged -->\n\nreal\n\n<!-- agentteam-runtime-input:real -->"}
	}`))
	events := drainKimiParsedEvents(t, driver, 2)
	require.Equal(t, "real", events[1].RuntimeInputID)

	driver.handleNotification("event", json.RawMessage(`{
		"type":"SteerInput",
		"payload":{"user_input":"<!-- agentteam-runtime-input:not-trailing -->\n\nmore text"}
	}`))
	select {
	case line := <-driver.lines:
		t.Fatalf("non-trailing marker acked: %s", line)
	default:
	}
}

func TestKimiInterruptActiveRunUsesCancelRequest(t *testing.T) {
	d := newKimiTestDriver(2)
	d.turnActive = true
	d.turnTerminal = make(chan struct{})
	var method string
	d.requestFuncForTest = func(_ context.Context, got string, _ any) (json.RawMessage, error) {
		method = got
		d.mu.Lock()
		d.markTurnTerminalLocked()
		d.mu.Unlock()
		return json.RawMessage(`{}`), nil
	}
	require.NoError(t, d.InterruptActiveRun(context.Background(), kimiTestSessionID))
	require.Equal(t, "cancel", method)
}

func TestKimiInterruptTreatsConcurrentTerminalAsSuccess(t *testing.T) {
	d := newKimiTestDriver(2)
	d.turnActive = true
	d.turnTerminal = make(chan struct{})
	d.requestFuncForTest = func(_ context.Context, got string, _ any) (json.RawMessage, error) {
		require.Equal(t, "cancel", got)
		d.mu.Lock()
		d.markTurnTerminalLocked()
		d.mu.Unlock()
		return nil, &kimiWireResponseError{Code: -32602, Message: "No agent turn is in progress"}
	}
	require.NoError(t, d.InterruptActiveRun(context.Background(), kimiTestSessionID))
}

func TestKimiInterruptingInputCancelsPromptBeforeTurnBegin(t *testing.T) {
	d := newKimiTestDriver(2)
	d.promptInFlight = true
	d.turnTerminal = make(chan struct{})
	requests := make(chan string, 2)
	d.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		requests <- method
		switch method {
		case "cancel":
			d.mu.Lock()
			d.promptInFlight = false
			d.markTurnTerminalLocked()
			d.mu.Unlock()
			return json.RawMessage(`{}`), nil
		case "prompt":
			d.mu.Lock()
			d.turnActive = true
			accepted := d.promptAccepted
			d.promptAccepted = nil
			d.mu.Unlock()
			accepted <- nil
			return json.RawMessage(`{"status":"finished"}`), nil
		default:
			t.Fatalf("unexpected method %q", method)
			return nil, nil
		}
	}

	require.NoError(t, d.DeliverInterruptingRuntimeInput(
		context.Background(),
		testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "replacement"}, "replacement"),
		kimiTestSessionID,
	))
	require.Equal(t, "cancel", <-requests)
	require.Equal(t, "prompt", <-requests)
}

func TestKimiReconcileRuntimeInputUsesDurableMarker(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	wire := strings.Join([]string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "hello\n\n" + kimiRuntimeInputMarker("acked")}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
	}, "\n") + "\n"
	require.NoError(t, os.WriteFile(filepath.Join(sessionDir, "wire.jsonl"), []byte(wire), 0o600))
	d := &KimiWireDriver{launchCfg: drivercontract.LaunchConfig{
		WorkDir: workDir,
		Env:     map[string]string{"KIMI_SHARE_DIR": shareDir},
	}}
	got, err := d.ReconcileRuntimeInput(context.Background(), "acked", kimiTestSessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationAcked, got)
	got, err = d.ReconcileRuntimeInput(context.Background(), "absent", kimiTestSessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationAbsent, got)
}

func TestKimiReconcileRuntimeInputIgnoresAssistantAndToolMarkers(t *testing.T) {
	workDir := t.TempDir()
	shareDir := t.TempDir()
	sessionDir := kimiTestSessionDir(t, shareDir, workDir, kimiTestSessionID)
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	writeKimiLines(t, filepath.Join(sessionDir, "wire.jsonl"), []string{
		`{"type":"metadata","protocol_version":"1.10"}`,
		kimiTestWireRecord("TurnBegin", map[string]any{"user_input": "hello"}),
		kimiTestWireRecord("ContentPart", map[string]any{"type": "text", "text": kimiRuntimeInputMarker("forged")}),
		kimiTestWireRecord("ToolResult", map[string]any{"output": kimiRuntimeInputMarker("forged")}),
		kimiTestWireRecord("TurnEnd", map[string]any{}),
	})
	d := &KimiWireDriver{launchCfg: drivercontract.LaunchConfig{
		WorkDir: workDir,
		Env:     map[string]string{"KIMI_SHARE_DIR": shareDir},
	}}
	got, err := d.ReconcileRuntimeInput(context.Background(), "forged", kimiTestSessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationAbsent, got)
}

func TestNewDriverKimiUsesACPDriver(t *testing.T) {
	driver, err := NewDriver(runtimeconfig.TypeKimi)
	require.NoError(t, err)
	_, ok := driver.(*KimiACPDriver)
	require.True(t, ok, "NewDriver(TypeKimi) returned %T", driver)
}

func kimiTestSessionDir(t *testing.T, shareDir, workDir, sessionID string) string {
	t.Helper()
	path, err := kimiSessionDir(drivercontract.LaunchConfig{
		WorkDir: workDir,
		Env:     map[string]string{"KIMI_SHARE_DIR": shareDir},
	}, sessionID, false)
	require.NoError(t, err)
	return path
}

func kimiTestWireRecord(messageType string, payload any) string {
	data, _ := json.Marshal(map[string]any{
		"timestamp": 1,
		"message": map[string]any{
			"type":    messageType,
			"payload": payload,
		},
	})
	return string(data)
}

type nopWriteCloser struct {
	Writer *bytes.Buffer
}

func (n nopWriteCloser) Write(p []byte) (int, error) {
	if n.Writer == nil {
		return len(p), nil
	}
	return n.Writer.Write(p)
}
func (nopWriteCloser) Close() error { return nil }

func newKimiTestDriver(buffer int) *KimiWireDriver {
	output := newKimiTranslatedOutput(buffer)
	return &KimiWireDriver{
		lines:        output.lines,
		output:       output,
		sessionID:    kimiTestSessionID,
		pending:      make(map[string]chan kimiWireEnvelope),
		toolNames:    make(map[string]string),
		turnTerminal: closedKimiSignal(),
	}
}

func argumentIndex(t *testing.T, args []string, argument string) int {
	t.Helper()
	for index, value := range args {
		if value == argument {
			return index
		}
	}
	t.Fatalf("argument %q not found in %#v", argument, args)
	return -1
}

func drainKimiParsedEvents(t *testing.T, d *KimiWireDriver, count int) []drivercontract.ParsedEvent {
	t.Helper()
	var events []drivercontract.ParsedEvent
	for len(events) < count {
		select {
		case line := <-d.lines:
			events = append(events, d.ParseLine(line)...)
		case <-time.After(time.Second):
			t.Fatalf("timed out waiting for Kimi event %d/%d", len(events), count)
		}
	}
	return events
}

func drainKimiOutput(t *testing.T, d *KimiWireDriver) []drivercontract.ParsedEvent {
	t.Helper()
	var events []drivercontract.ParsedEvent
	for {
		select {
		case line := <-d.output.lines:
			events = append(events, d.ParseLine(line)...)
		default:
			return events
		}
	}
}

func TestKimiStatusUpdateDeduplicatesByMessageID(t *testing.T) {
	t.Parallel()

	output := newKimiTranslatedOutput(4)
	d := &KimiWireDriver{output: output, sessionID: "sess-1"}

	payload := json.RawMessage(`{"message_id":"res-1","token_usage":{"input_other":100,"input_cache_read":20,"input_cache_creation":10,"output":50}}`)

	d.handleKimiStatusUpdate(payload, "sess-1")
	d.handleKimiStatusUpdate(payload, "sess-1") // replayed StatusUpdate must be deduped

	events := drainKimiOutput(t, d)
	var usage []drivercontract.ParsedEvent
	for _, ev := range events {
		if ev.Kind == consts.ParsedEventKindUsage {
			usage = append(usage, ev)
		}
	}
	require.Len(t, usage, 1, "replayed StatusUpdate should not emit duplicate usage")
	require.Equal(t, "sess-1", usage[0].SessionID)
	require.Equal(t, "res-1", usage[0].ModelInvocationID)
	require.True(t, usage[0].ContextExposureKnown)
	require.False(t, usage[0].ContextExposureExcluded)
	require.Equal(t, uint64(130), usage[0].InputTokens) // 100 + 20 + 10
	require.Equal(t, uint64(20), usage[0].CachedInputTokens)
	require.Equal(t, uint64(50), usage[0].OutputTokens)
}

func TestKimiStatusUpdateWithoutMessageIDIsNotDeduped(t *testing.T) {
	t.Parallel()

	output := newKimiTranslatedOutput(4)
	d := &KimiWireDriver{output: output, sessionID: "sess-1"}

	// Without a message_id the driver cannot dedup, so each call emits (preserves
	// prior behavior for events lacking the stable id).
	payload := json.RawMessage(`{"token_usage":{"input_other":10,"input_cache_read":0,"input_cache_creation":0,"output":5}}`)
	d.handleKimiStatusUpdate(payload, "sess-1")
	d.handleKimiStatusUpdate(payload, "sess-1")

	events := drainKimiOutput(t, d)
	count := 0
	for _, ev := range events {
		if ev.Kind == consts.ParsedEventKindUsage {
			require.Empty(t, ev.ModelInvocationID)
			count++
		}
	}
	require.Equal(t, 2, count)
}

func TestKimiSubagentEventCountsEmbeddedStatusUsage(t *testing.T) {
	t.Parallel()

	output := newKimiTranslatedOutput(4)
	d := &KimiWireDriver{output: output, sessionID: "sess-1"}

	// A subagent's per-step StatusUpdate is embedded inside a SubagentEvent; its
	// tokens must be counted rather than dropped.
	subagentPayload := json.RawMessage(`{
		"parent_tool_call_id":"call-1","agent_id":"agent-sub","subagent_type":"researcher",
		"event":{"type":"StatusUpdate","payload":{"message_id":"sub-res-1","token_usage":{"input_other":40,"input_cache_read":5,"input_cache_creation":0,"output":15}}}
	}`)
	d.handleKimiSubagentEvent(subagentPayload, "sess-1")

	var usage *drivercontract.ParsedEvent
	for _, ev := range drainKimiOutput(t, d) {
		if ev.Kind == consts.ParsedEventKindUsage {
			cp := ev
			usage = &cp
		}
	}
	require.NotNil(t, usage, "subagent StatusUpdate usage was not counted")
	require.False(t, usage.ContextExposureKnown)
	require.True(t, usage.ContextExposureExcluded)
	require.Equal(t, uint64(45), usage.InputTokens) // 40 + 5 + 0
	require.Equal(t, uint64(5), usage.CachedInputTokens)
	require.Equal(t, uint64(15), usage.OutputTokens)
}