初始化仓库
This commit is contained in:
@@ -0,0 +1,461 @@
|
||||
package kimi
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
)
|
||||
|
||||
func init() {
|
||||
core.RegisterAgent("kimi", New)
|
||||
}
|
||||
|
||||
// Agent drives Kimi Code CLI using --print --output-format stream-json.
|
||||
//
|
||||
// Modes:
|
||||
// - "default": standard mode (note: --print implicitly enables --yolo)
|
||||
// - "yolo": auto-approve all tool calls
|
||||
// - "plan": read-only plan mode
|
||||
// - "quiet": alias for --quiet (print + text + final-message-only)
|
||||
type Agent struct {
|
||||
workDir string
|
||||
model string
|
||||
mode string
|
||||
cmd string // CLI binary name, default "kimi"
|
||||
timeout time.Duration
|
||||
providers []core.ProviderConfig
|
||||
activeIdx int // -1 = no provider set
|
||||
sessionEnv []string
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func New(opts map[string]any) (core.Agent, error) {
|
||||
workDir, _ := opts["work_dir"].(string)
|
||||
if workDir == "" {
|
||||
workDir = "."
|
||||
}
|
||||
model, _ := opts["model"].(string)
|
||||
mode, _ := opts["mode"].(string)
|
||||
mode = normalizeMode(mode)
|
||||
cmd, _ := opts["cmd"].(string)
|
||||
if cmd == "" {
|
||||
cmd = "kimi"
|
||||
}
|
||||
|
||||
var timeoutMins int64
|
||||
switch v := opts["timeout_mins"].(type) {
|
||||
case int64:
|
||||
timeoutMins = v
|
||||
case int:
|
||||
timeoutMins = int64(v)
|
||||
case float64:
|
||||
timeoutMins = int64(v)
|
||||
default:
|
||||
if v != nil {
|
||||
slog.Debug("kimi: timeout_mins has unexpected type", "type", fmt.Sprintf("%T", v))
|
||||
}
|
||||
}
|
||||
var timeout time.Duration
|
||||
if timeoutMins > 0 {
|
||||
timeout = time.Duration(timeoutMins) * time.Minute
|
||||
}
|
||||
|
||||
if _, err := exec.LookPath(cmd); err != nil {
|
||||
return nil, fmt.Errorf("kimi: %q CLI not found in PATH, install with: pip install kimi-cli", cmd)
|
||||
}
|
||||
|
||||
return &Agent{
|
||||
workDir: workDir,
|
||||
model: model,
|
||||
mode: mode,
|
||||
cmd: cmd,
|
||||
timeout: timeout,
|
||||
activeIdx: -1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func normalizeMode(raw string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(raw)) {
|
||||
case "yolo", "force", "bypass", "auto":
|
||||
return "yolo"
|
||||
case "plan":
|
||||
return "plan"
|
||||
case "quiet":
|
||||
return "quiet"
|
||||
default:
|
||||
return "default"
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) Name() string { return "kimi" }
|
||||
func (a *Agent) CLIBinaryName() string { return "kimi" }
|
||||
func (a *Agent) CLIDisplayName() string { return "Kimi" }
|
||||
|
||||
func (a *Agent) SetWorkDir(dir string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.workDir = dir
|
||||
slog.Info("kimi: work_dir changed", "work_dir", dir)
|
||||
}
|
||||
|
||||
func (a *Agent) GetWorkDir() string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return a.workDir
|
||||
}
|
||||
|
||||
func (a *Agent) SetModel(model string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.model = model
|
||||
slog.Info("kimi: model changed", "model", model)
|
||||
}
|
||||
|
||||
func (a *Agent) GetModel() string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return core.GetProviderModel(a.providers, a.activeIdx, a.model)
|
||||
}
|
||||
|
||||
func (a *Agent) configuredModels() []core.ModelOption {
|
||||
a.mu.RLock()
|
||||
defer a.mu.RUnlock()
|
||||
return core.GetProviderModels(a.providers, a.activeIdx)
|
||||
}
|
||||
|
||||
func (a *Agent) AvailableModels(ctx context.Context) []core.ModelOption {
|
||||
if models := a.configuredModels(); len(models) > 0 {
|
||||
return models
|
||||
}
|
||||
return []core.ModelOption{
|
||||
{Name: "kimi-k2-0711-preview", Desc: "Kimi K2 (most capable)"},
|
||||
{Name: "kimi-k2-0711", Desc: "Kimi K2"},
|
||||
{Name: "kimi-k2-5-preview", Desc: "Kimi K2.5 (balanced)"},
|
||||
{Name: "kimi-k2-5", Desc: "Kimi K2.5"},
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) SetSessionEnv(env []string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.sessionEnv = env
|
||||
}
|
||||
|
||||
func (a *Agent) StartSession(ctx context.Context, sessionID string) (core.AgentSession, error) {
|
||||
a.mu.Lock()
|
||||
model := a.model
|
||||
mode := a.mode
|
||||
cmd := a.cmd
|
||||
workDir := a.workDir
|
||||
timeout := a.timeout
|
||||
extraEnv := a.providerEnvLocked()
|
||||
extraEnv = append(extraEnv, a.sessionEnv...)
|
||||
if a.activeIdx >= 0 && a.activeIdx < len(a.providers) {
|
||||
if m := a.providers[a.activeIdx].Model; m != "" {
|
||||
model = m
|
||||
}
|
||||
}
|
||||
a.mu.Unlock()
|
||||
|
||||
return newKimiSession(ctx, cmd, workDir, model, mode, sessionID, extraEnv, timeout)
|
||||
}
|
||||
|
||||
func (a *Agent) ListSessions(_ context.Context) ([]core.AgentSessionInfo, error) {
|
||||
return listKimiSessions(a.workDir)
|
||||
}
|
||||
|
||||
func (a *Agent) DeleteSession(_ context.Context, sessionID string) error {
|
||||
path := findKimiSessionDir(sessionID)
|
||||
if path == "" {
|
||||
return fmt.Errorf("session not found: %s", sessionID)
|
||||
}
|
||||
return os.RemoveAll(path)
|
||||
}
|
||||
|
||||
func (a *Agent) Stop() error { return nil }
|
||||
|
||||
// ── ModeSwitcher ────────────────────────────────────────────────
|
||||
|
||||
func (a *Agent) SetMode(mode string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.mode = normalizeMode(mode)
|
||||
slog.Info("kimi: mode changed", "mode", a.mode)
|
||||
}
|
||||
|
||||
func (a *Agent) GetMode() string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return a.mode
|
||||
}
|
||||
|
||||
func (a *Agent) PermissionModes() []core.PermissionModeInfo {
|
||||
return []core.PermissionModeInfo{
|
||||
{Key: "default", Name: "Default", NameZh: "默认", Desc: "Standard mode (print output)", DescZh: "标准模式（打印输出）"},
|
||||
{Key: "yolo", Name: "YOLO", NameZh: "全自动", Desc: "Auto-approve all tool calls", DescZh: "自动批准所有工具调用"},
|
||||
{Key: "plan", Name: "Plan", NameZh: "规划模式", Desc: "Read-only plan mode, no execution", DescZh: "只读规划模式，不做修改"},
|
||||
{Key: "quiet", Name: "Quiet", NameZh: "静默", Desc: "Quiet mode (final message only)", DescZh: "静默模式（仅最终消息）"},
|
||||
}
|
||||
}
|
||||
|
||||
// ── SkillProvider implementation ──────────────────────────────
|
||||
|
||||
func (a *Agent) SkillDirs() []string {
|
||||
absDir, err := filepath.Abs(a.workDir)
|
||||
if err != nil {
|
||||
absDir = a.workDir
|
||||
}
|
||||
dirs := []string{filepath.Join(absDir, ".kimi", "skills")}
|
||||
if home, err := os.UserHomeDir(); err == nil {
|
||||
dirs = append(dirs, filepath.Join(home, ".kimi", "skills"))
|
||||
}
|
||||
return dirs
|
||||
}
|
||||
|
||||
// ── ContextCompressor implementation ──────────────────────────
|
||||
|
||||
func (a *Agent) CompressCommand() string { return "" }
|
||||
|
||||
// ── MemoryFileProvider implementation ─────────────────────────
|
||||
|
||||
func (a *Agent) ProjectMemoryFile() string {
|
||||
absDir, err := filepath.Abs(a.workDir)
|
||||
if err != nil {
|
||||
absDir = a.workDir
|
||||
}
|
||||
return filepath.Join(absDir, "AGENTS.md")
|
||||
}
|
||||
|
||||
func (a *Agent) GlobalMemoryFile() string {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return filepath.Join(homeDir, ".kimi", "AGENTS.md")
|
||||
}
|
||||
|
||||
// ── ProviderSwitcher ────────────────────────────────────────────
|
||||
|
||||
func (a *Agent) SetProviders(providers []core.ProviderConfig) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.providers = providers
|
||||
}
|
||||
|
||||
func (a *Agent) SetActiveProvider(name string) bool {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if name == "" {
|
||||
a.activeIdx = -1
|
||||
slog.Info("kimi: provider cleared")
|
||||
return true
|
||||
}
|
||||
for i, p := range a.providers {
|
||||
if p.Name == name {
|
||||
a.activeIdx = i
|
||||
slog.Info("kimi: provider switched", "provider", name)
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (a *Agent) GetActiveProvider() *core.ProviderConfig {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if a.activeIdx < 0 || a.activeIdx >= len(a.providers) {
|
||||
return nil
|
||||
}
|
||||
p := a.providers[a.activeIdx]
|
||||
return &p
|
||||
}
|
||||
|
||||
func (a *Agent) ListProviders() []core.ProviderConfig {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
result := make([]core.ProviderConfig, len(a.providers))
|
||||
copy(result, a.providers)
|
||||
return result
|
||||
}
|
||||
|
||||
func (a *Agent) providerEnvLocked() []string {
|
||||
if a.activeIdx < 0 || a.activeIdx >= len(a.providers) {
|
||||
return nil
|
||||
}
|
||||
p := a.providers[a.activeIdx]
|
||||
var env []string
|
||||
if p.APIKey != "" {
|
||||
env = append(env, "KIMI_API_KEY="+p.APIKey)
|
||||
}
|
||||
for k, v := range p.Env {
|
||||
env = append(env, k+"="+v)
|
||||
}
|
||||
return env
|
||||
}
|
||||
|
||||
// ── Session listing ─────────────────────────────────────────────
|
||||
|
||||
func kimiSessionsBaseDir() string {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return filepath.Join(homeDir, ".kimi", "sessions")
|
||||
}
|
||||
|
||||
func listKimiSessions(workDir string) ([]core.AgentSessionInfo, error) {
|
||||
absWorkDir, err := filepath.Abs(workDir)
|
||||
if err != nil {
|
||||
absWorkDir = workDir
|
||||
}
|
||||
|
||||
sessionsBase := kimiSessionsBaseDir()
|
||||
if sessionsBase == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(sessionsBase)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("kimi: read sessions dir: %w", err)
|
||||
}
|
||||
|
||||
var sessions []core.AgentSessionInfo
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
projectDir := filepath.Join(sessionsBase, entry.Name())
|
||||
sessionEntries, err := os.ReadDir(projectDir)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, se := range sessionEntries {
|
||||
if !se.IsDir() {
|
||||
continue
|
||||
}
|
||||
sessionDir := filepath.Join(projectDir, se.Name())
|
||||
info := parseKimiSessionDir(sessionDir, absWorkDir)
|
||||
if info != nil {
|
||||
sessions = append(sessions, *info)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sort.Slice(sessions, func(i, j int) bool {
|
||||
return sessions[i].ModifiedAt.After(sessions[j].ModifiedAt)
|
||||
})
|
||||
|
||||
return sessions, nil
|
||||
}
|
||||
|
||||
func parseKimiSessionDir(sessionDir, filterWorkDir string) *core.AgentSessionInfo {
|
||||
statePath := filepath.Join(sessionDir, "state.json")
|
||||
stateData, err := os.ReadFile(statePath)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var state struct {
|
||||
CustomTitle string `json:"custom_title"`
|
||||
Archived bool `json:"archived"`
|
||||
}
|
||||
if json.Unmarshal(stateData, &state) != nil {
|
||||
return nil
|
||||
}
|
||||
if state.Archived {
|
||||
return nil
|
||||
}
|
||||
|
||||
sessionID := filepath.Base(sessionDir)
|
||||
|
||||
info, err := os.Stat(sessionDir)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
msgCount := 0
|
||||
summary := ""
|
||||
contextPath := filepath.Join(sessionDir, "context.jsonl")
|
||||
if f, err := os.Open(contextPath); err == nil {
|
||||
defer f.Close()
|
||||
scanner := bufio.NewScanner(f)
|
||||
scanner.Buffer(make([]byte, 256*1024), 256*1024)
|
||||
for scanner.Scan() {
|
||||
var entry struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if json.Unmarshal(scanner.Bytes(), &entry) != nil {
|
||||
continue
|
||||
}
|
||||
if entry.Role == "user" || entry.Role == "assistant" {
|
||||
msgCount++
|
||||
if entry.Role == "user" && entry.Content != "" && summary == "" {
|
||||
summary = strings.TrimSpace(entry.Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_ = filterWorkDir // Kimi does not store cwd in session metadata; list all sessions.
|
||||
|
||||
if summary == "" {
|
||||
summary = state.CustomTitle
|
||||
}
|
||||
if utf8.RuneCountInString(summary) > 60 {
|
||||
summary = string([]rune(summary)[:60]) + "..."
|
||||
}
|
||||
|
||||
return &core.AgentSessionInfo{
|
||||
ID: sessionID,
|
||||
Summary: summary,
|
||||
MessageCount: msgCount,
|
||||
ModifiedAt: info.ModTime(),
|
||||
}
|
||||
}
|
||||
|
||||
func findKimiSessionDir(sessionID string) string {
|
||||
sessionsBase := kimiSessionsBaseDir()
|
||||
if sessionsBase == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(sessionsBase)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
projectDir := filepath.Join(sessionsBase, entry.Name())
|
||||
sessionEntries, err := os.ReadDir(projectDir)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, se := range sessionEntries {
|
||||
if !se.IsDir() {
|
||||
continue
|
||||
}
|
||||
if se.Name() == sessionID {
|
||||
return filepath.Join(projectDir, se.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
package kimi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func skipUnlessKimiAvailable(t *testing.T) {
|
||||
t.Helper()
|
||||
if _, err := exec.LookPath("kimi"); err != nil {
|
||||
t.Skipf("kimi CLI not in PATH, skipping: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeMode(t *testing.T) {
|
||||
cases := []struct {
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{"default", "default"},
|
||||
{"DEFAULT", "default"},
|
||||
{"yolo", "yolo"},
|
||||
{"YOLO", "yolo"},
|
||||
{"force", "yolo"},
|
||||
{"bypass", "yolo"},
|
||||
{"auto", "yolo"},
|
||||
{"plan", "plan"},
|
||||
{"quiet", "quiet"},
|
||||
{"", "default"},
|
||||
{"unknown", "default"},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
assert.Equal(t, c.expected, normalizeMode(c.input), "input: %s", c.input)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentNew(t *testing.T) {
|
||||
skipUnlessKimiAvailable(t)
|
||||
agentInf, err := New(map[string]any{
|
||||
"work_dir": "/tmp",
|
||||
"model": "kimi-k2",
|
||||
"mode": "yolo",
|
||||
"timeout_mins": 15,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, agentInf)
|
||||
|
||||
a := agentInf.(*Agent)
|
||||
assert.Equal(t, "kimi", a.Name())
|
||||
assert.Equal(t, "/tmp", a.GetWorkDir())
|
||||
assert.Equal(t, "yolo", a.GetMode())
|
||||
assert.Equal(t, "kimi-k2", a.GetModel())
|
||||
}
|
||||
|
||||
// TestAgentFields verifies Name/WorkDir/Mode/Model without requiring
|
||||
// the kimi CLI on PATH — constructs the struct directly.
|
||||
func TestAgentFields(t *testing.T) {
|
||||
a := &Agent{
|
||||
workDir: "/tmp",
|
||||
model: "kimi-k2",
|
||||
mode: "yolo",
|
||||
cmd: "kimi",
|
||||
activeIdx: -1,
|
||||
}
|
||||
assert.Equal(t, "kimi", a.Name())
|
||||
assert.Equal(t, "Kimi", a.CLIDisplayName())
|
||||
assert.Equal(t, "kimi", a.CLIBinaryName())
|
||||
assert.Equal(t, "/tmp", a.GetWorkDir())
|
||||
assert.Equal(t, "yolo", a.GetMode())
|
||||
assert.Equal(t, "kimi-k2", a.GetModel())
|
||||
}
|
||||
|
||||
func TestAgentSetters(t *testing.T) {
|
||||
a := &Agent{workDir: "/tmp", mode: "default", activeIdx: -1}
|
||||
|
||||
a.SetWorkDir("/new/path")
|
||||
assert.Equal(t, "/new/path", a.GetWorkDir())
|
||||
|
||||
a.SetModel("kimi-k2-5")
|
||||
assert.Equal(t, "kimi-k2-5", a.GetModel())
|
||||
|
||||
a.SetMode("plan")
|
||||
assert.Equal(t, "plan", a.GetMode())
|
||||
}
|
||||
|
||||
func TestAgentPermissionModes(t *testing.T) {
|
||||
a := &Agent{}
|
||||
|
||||
modes := a.PermissionModes()
|
||||
require.Len(t, modes, 4)
|
||||
assert.Equal(t, "default", modes[0].Key)
|
||||
assert.Equal(t, "yolo", modes[1].Key)
|
||||
assert.Equal(t, "plan", modes[2].Key)
|
||||
assert.Equal(t, "quiet", modes[3].Key)
|
||||
}
|
||||
|
||||
func TestAgentProviderSwitcher(t *testing.T) {
|
||||
a := &Agent{workDir: "/tmp", activeIdx: -1}
|
||||
|
||||
providers := []core.ProviderConfig{
|
||||
{Name: "moonshot", APIKey: "sk-123"},
|
||||
{Name: "custom", BaseURL: "https://api.example.com"},
|
||||
}
|
||||
a.SetProviders(providers)
|
||||
|
||||
assert.False(t, a.SetActiveProvider("missing"))
|
||||
assert.True(t, a.SetActiveProvider("moonshot"))
|
||||
assert.Equal(t, "moonshot", a.GetActiveProvider().Name)
|
||||
|
||||
list := a.ListProviders()
|
||||
require.Len(t, list, 2)
|
||||
assert.Equal(t, "moonshot", list[0].Name)
|
||||
}
|
||||
|
||||
func TestAgentStartSession(t *testing.T) {
|
||||
skipUnlessKimiAvailable(t)
|
||||
agentInf, err := New(map[string]any{
|
||||
"work_dir": "/tmp",
|
||||
"model": "kimi-k2",
|
||||
"mode": "default",
|
||||
"timeout_mins": 10,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
session, err := agentInf.StartSession(ctx, "test-session-id")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, session)
|
||||
assert.True(t, session.Alive())
|
||||
assert.Equal(t, "test-session-id", session.CurrentSessionID())
|
||||
|
||||
err = session.Close()
|
||||
assert.NoError(t, err)
|
||||
assert.False(t, session.Alive())
|
||||
}
|
||||
|
||||
func TestAgentMemoryAndSkill(t *testing.T) {
|
||||
a := &Agent{workDir: "/tmp/my-project", activeIdx: -1}
|
||||
|
||||
assert.Equal(t, "/tmp/my-project/AGENTS.md", a.ProjectMemoryFile())
|
||||
assert.NotEmpty(t, a.GlobalMemoryFile())
|
||||
|
||||
skillDirs := a.SkillDirs()
|
||||
require.Len(t, skillDirs, 2)
|
||||
assert.Contains(t, skillDirs[0], ".kimi/skills")
|
||||
assert.Contains(t, skillDirs[1], ".kimi/skills")
|
||||
}
|
||||
|
||||
func TestAgentAvailableModels(t *testing.T) {
|
||||
a := &Agent{workDir: "/tmp", activeIdx: -1}
|
||||
|
||||
models := a.AvailableModels(context.Background())
|
||||
require.True(t, len(models) > 0)
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
package kimi
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
)
|
||||
|
||||
// kimSession manages multi-turn conversations with the Kimi CLI.
|
||||
// Each Send() launches a new `kimi --print --output-format stream-json` process
|
||||
// with --resume for conversation continuity.
|
||||
type kimiSession struct {
|
||||
cmd string
|
||||
workDir string
|
||||
model string
|
||||
mode string
|
||||
timeout time.Duration
|
||||
extraEnv []string
|
||||
events chan core.Event
|
||||
sessionID atomic.Value // stores string — Kimi session ID
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
alive atomic.Bool
|
||||
|
||||
pendingMsgs []string // buffered assistant text messages
|
||||
}
|
||||
|
||||
func newKimiSession(ctx context.Context, cmd, workDir, model, mode, resumeID string, extraEnv []string, timeout time.Duration) (*kimiSession, error) {
|
||||
sessionCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
ks := &kimiSession{
|
||||
cmd: cmd,
|
||||
workDir: workDir,
|
||||
model: model,
|
||||
mode: mode,
|
||||
timeout: timeout,
|
||||
extraEnv: extraEnv,
|
||||
events: make(chan core.Event, 64),
|
||||
ctx: sessionCtx,
|
||||
cancel: cancel,
|
||||
}
|
||||
ks.alive.Store(true)
|
||||
|
||||
if resumeID != "" && resumeID != core.ContinueSession {
|
||||
ks.sessionID.Store(resumeID)
|
||||
}
|
||||
|
||||
return ks, nil
|
||||
}
|
||||
|
||||
func (ks *kimiSession) Send(prompt string, images []core.ImageAttachment, files []core.FileAttachment) error {
|
||||
if !ks.alive.Load() {
|
||||
return fmt.Errorf("session is closed")
|
||||
}
|
||||
|
||||
// Save images and files into the workspace so Kimi CLI can access them.
|
||||
attachDir := filepath.Join(ks.workDir, ".cc-connect", "attachments")
|
||||
if (len(images) > 0 || len(files) > 0) && os.MkdirAll(attachDir, 0o755) != nil {
|
||||
attachDir = os.TempDir()
|
||||
}
|
||||
|
||||
var imageRefs []string
|
||||
for i, img := range images {
|
||||
ext := ".png"
|
||||
switch img.MimeType {
|
||||
case "image/jpeg":
|
||||
ext = ".jpg"
|
||||
case "image/gif":
|
||||
ext = ".gif"
|
||||
case "image/webp":
|
||||
ext = ".webp"
|
||||
}
|
||||
fname := fmt.Sprintf("img_%d_%d%s", time.Now().UnixMilli(), i, ext)
|
||||
fpath := filepath.Join(attachDir, fname)
|
||||
if err := os.WriteFile(fpath, img.Data, 0o644); err != nil {
|
||||
slog.Warn("kimiSession: failed to save image", "error", err)
|
||||
continue
|
||||
}
|
||||
imageRefs = append(imageRefs, fpath)
|
||||
}
|
||||
|
||||
var fileRefs []string
|
||||
for i, f := range files {
|
||||
fname := filepath.Base(f.FileName)
|
||||
if fname == "" || fname == "." || fname == ".." {
|
||||
fname = fmt.Sprintf("file_%d_%d", time.Now().UnixMilli(), i)
|
||||
}
|
||||
fpath := filepath.Join(attachDir, fname)
|
||||
if err := os.WriteFile(fpath, f.Data, 0o644); err != nil {
|
||||
slog.Warn("kimiSession: failed to save file", "error", err)
|
||||
continue
|
||||
}
|
||||
fileRefs = append(fileRefs, fpath)
|
||||
}
|
||||
|
||||
fullPrompt := prompt
|
||||
if len(imageRefs) > 0 {
|
||||
if fullPrompt == "" {
|
||||
fullPrompt = "Please analyze the attached image(s)."
|
||||
}
|
||||
fullPrompt += "\n\n[Attached images saved at: " + strings.Join(imageRefs, ", ") + "]"
|
||||
}
|
||||
if len(fileRefs) > 0 {
|
||||
if fullPrompt == "" {
|
||||
fullPrompt = "Please analyze the attached file(s)."
|
||||
}
|
||||
fullPrompt += "\n\n[Attached files saved at: " + strings.Join(fileRefs, ", ") + "]"
|
||||
}
|
||||
|
||||
args := []string{
|
||||
"--print",
|
||||
"--output-format", "stream-json",
|
||||
}
|
||||
|
||||
switch ks.mode {
|
||||
case "plan":
|
||||
args = append(args, "--plan")
|
||||
case "quiet":
|
||||
args = append(args, "--quiet")
|
||||
}
|
||||
|
||||
sid := ks.CurrentSessionID()
|
||||
if sid != "" {
|
||||
args = append(args, "--resume", sid)
|
||||
}
|
||||
if ks.model != "" {
|
||||
args = append(args, "--model", ks.model)
|
||||
}
|
||||
if ks.workDir != "" {
|
||||
args = append(args, "--work-dir", ks.workDir)
|
||||
}
|
||||
|
||||
args = append(args, "--prompt", fullPrompt)
|
||||
|
||||
var cancel context.CancelFunc
|
||||
var ctx context.Context
|
||||
if ks.timeout > 0 {
|
||||
ctx, cancel = context.WithTimeout(ks.ctx, ks.timeout)
|
||||
} else {
|
||||
ctx, cancel = context.WithCancel(ks.ctx)
|
||||
}
|
||||
|
||||
started := false
|
||||
defer func() {
|
||||
if !started {
|
||||
cancel()
|
||||
}
|
||||
}()
|
||||
|
||||
slog.Debug("kimiSession: launching", "resume", sid != "", "args", core.RedactArgs(args))
|
||||
cmd := exec.CommandContext(ctx, ks.cmd, args...)
|
||||
cmd.WaitDelay = 1 * time.Second
|
||||
cmd.Dir = ks.workDir
|
||||
env := os.Environ()
|
||||
if len(ks.extraEnv) > 0 {
|
||||
env = core.MergeEnv(env, ks.extraEnv)
|
||||
}
|
||||
cmd.Env = env
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return fmt.Errorf("kimiSession: stdout pipe: %w", err)
|
||||
}
|
||||
|
||||
var stderrBuf bytes.Buffer
|
||||
cmd.Stderr = &stderrBuf
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return fmt.Errorf("kimiSession: start: %w", err)
|
||||
}
|
||||
|
||||
started = true
|
||||
ks.wg.Add(1)
|
||||
go func() {
|
||||
defer cancel()
|
||||
ks.readLoop(ctx, cmd, stdout, &stderrBuf, append(imageRefs, fileRefs...))
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ks *kimiSession) readLoop(ctx context.Context, cmd *exec.Cmd, stdout io.ReadCloser, stderrBuf *bytes.Buffer, tempFiles []string) {
|
||||
defer ks.wg.Done()
|
||||
defer func() {
|
||||
for _, f := range tempFiles {
|
||||
os.Remove(f)
|
||||
}
|
||||
}()
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
stdout.Close()
|
||||
}()
|
||||
|
||||
scanner := bufio.NewScanner(stdout)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024)
|
||||
|
||||
var scanErr error
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
slog.Debug("kimiSession: raw", "line", truncate(line, 500))
|
||||
|
||||
// Kimi prints a non-JSON line at the end: "To resume this session: kimi -r <id>"
|
||||
if strings.HasPrefix(line, "To resume this session:") {
|
||||
if id := extractResumeSessionID(line); id != "" {
|
||||
ks.sessionID.Store(id)
|
||||
slog.Debug("kimiSession: session id updated", "session_id", id)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
var raw map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &raw); err != nil {
|
||||
slog.Debug("kimiSession: non-JSON line", "line", line)
|
||||
continue
|
||||
}
|
||||
|
||||
ks.handleEvent(raw)
|
||||
}
|
||||
scanErr = scanner.Err()
|
||||
|
||||
// Wait for process exit before sending any terminal event so the engine
|
||||
// never sees EventError after EventResult from the same turn.
|
||||
waitErr := cmd.Wait()
|
||||
|
||||
// Kimi writes "To resume this session: kimi -r <uuid>" to stderr (not stdout),
|
||||
// so the scanner above never sees it. Extract it from the captured stderr
|
||||
// buffer before emitting EventResult so the next turn can pass --resume.
|
||||
for _, line := range strings.Split(stderrBuf.String(), "\n") {
|
||||
if strings.HasPrefix(strings.TrimSpace(line), "To resume this session:") {
|
||||
if id := extractResumeSessionID(line); id != "" {
|
||||
ks.sessionID.Store(id)
|
||||
slog.Debug("kimiSession: session id from stderr", "session_id", id)
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if scanErr != nil {
|
||||
slog.Error("kimiSession: scanner error", "error", scanErr)
|
||||
evt := core.Event{Type: core.EventError, Error: fmt.Errorf("read stdout: %w", scanErr)}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if waitErr != nil {
|
||||
stderrMsg := strings.TrimSpace(stderrBuf.String())
|
||||
if stderrMsg != "" {
|
||||
slog.Error("kimiSession: process failed", "error", waitErr, "stderr", stderrMsg)
|
||||
evt := core.Event{Type: core.EventError, Error: fmt.Errorf("%s", stderrMsg)}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Flush any remaining pending messages as text and send result event.
|
||||
ks.flushPendingAsText()
|
||||
evt := core.Event{Type: core.EventResult, SessionID: ks.CurrentSessionID(), Done: true}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
func extractResumeSessionID(line string) string {
|
||||
// Format: "To resume this session: kimi -r <uuid>"
|
||||
parts := strings.Fields(line)
|
||||
for i, p := range parts {
|
||||
if p == "-r" && i+1 < len(parts) {
|
||||
return parts[i+1]
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// Kimi CLI stream-json message roles:
|
||||
// - "assistant": content (think + text), tool_calls
|
||||
// - "tool": content (tool execution result), tool_call_id
|
||||
func (ks *kimiSession) handleEvent(raw map[string]any) {
|
||||
role, _ := raw["role"].(string)
|
||||
|
||||
switch role {
|
||||
case "assistant":
|
||||
ks.handleAssistant(raw)
|
||||
case "tool":
|
||||
ks.handleTool(raw)
|
||||
default:
|
||||
slog.Debug("kimiSession: unhandled role", "role", role)
|
||||
}
|
||||
}
|
||||
|
||||
func (ks *kimiSession) handleAssistant(raw map[string]any) {
|
||||
content, _ := raw["content"].([]any)
|
||||
for _, item := range content {
|
||||
block, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
blockType, _ := block["type"].(string)
|
||||
switch blockType {
|
||||
case "think", "thinking":
|
||||
if think, ok := block["think"].(string); ok && think != "" {
|
||||
evt := core.Event{Type: core.EventThinking, Content: think}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
case "text":
|
||||
if text, ok := block["text"].(string); ok && text != "" {
|
||||
ks.pendingMsgs = append(ks.pendingMsgs, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Handle tool_calls
|
||||
toolCalls, _ := raw["tool_calls"].([]any)
|
||||
if len(toolCalls) > 0 {
|
||||
ks.flushPendingAsThinking()
|
||||
for _, tc := range toolCalls {
|
||||
tcMap, ok := tc.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
funcBlock, _ := tcMap["function"].(map[string]any)
|
||||
toolName, _ := funcBlock["name"].(string)
|
||||
args, _ := funcBlock["arguments"].(string)
|
||||
toolID, _ := tcMap["id"].(string)
|
||||
|
||||
slog.Debug("kimiSession: tool_call", "tool", toolName, "id", toolID)
|
||||
evt := core.Event{
|
||||
Type: core.EventToolUse,
|
||||
ToolName: toolName,
|
||||
ToolInput: truncate(strings.TrimSpace(args), 500),
|
||||
RequestID: toolID,
|
||||
}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ks *kimiSession) handleTool(raw map[string]any) {
|
||||
toolCallID, _ := raw["tool_call_id"].(string)
|
||||
content, _ := raw["content"].([]any)
|
||||
var outputParts []string
|
||||
for _, item := range content {
|
||||
block, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
blockType, _ := block["type"].(string)
|
||||
if blockType == "text" {
|
||||
if text, ok := block["text"].(string); ok {
|
||||
outputParts = append(outputParts, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
output := strings.Join(outputParts, "")
|
||||
|
||||
if output != "" {
|
||||
slog.Debug("kimiSession: tool result", "tool_call_id", toolCallID)
|
||||
evt := core.Event{
|
||||
Type: core.EventToolResult,
|
||||
ToolName: toolCallID,
|
||||
ToolResult: truncate(strings.TrimSpace(output), 500),
|
||||
}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ks *kimiSession) flushPendingAsThinking() {
|
||||
if len(ks.pendingMsgs) == 0 {
|
||||
return
|
||||
}
|
||||
text := strings.Join(ks.pendingMsgs, "")
|
||||
ks.pendingMsgs = ks.pendingMsgs[:0]
|
||||
if text != "" {
|
||||
evt := core.Event{Type: core.EventThinking, Content: text}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ks *kimiSession) flushPendingAsText() {
|
||||
if len(ks.pendingMsgs) == 0 {
|
||||
return
|
||||
}
|
||||
text := strings.Join(ks.pendingMsgs, "")
|
||||
ks.pendingMsgs = ks.pendingMsgs[:0]
|
||||
if text != "" {
|
||||
evt := core.Event{Type: core.EventText, Content: text}
|
||||
select {
|
||||
case ks.events <- evt:
|
||||
case <-ks.ctx.Done():
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// RespondPermission is a no-op — Kimi CLI permissions are handled via --print (implicit --yolo).
|
||||
func (ks *kimiSession) RespondPermission(_ string, _ core.PermissionResult) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ks *kimiSession) Events() <-chan core.Event {
|
||||
return ks.events
|
||||
}
|
||||
|
||||
func (ks *kimiSession) CurrentSessionID() string {
|
||||
v, _ := ks.sessionID.Load().(string)
|
||||
return v
|
||||
}
|
||||
|
||||
func (ks *kimiSession) Alive() bool {
|
||||
return ks.alive.Load()
|
||||
}
|
||||
|
||||
func (ks *kimiSession) Close() error {
|
||||
ks.alive.Store(false)
|
||||
ks.cancel()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
ks.wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
close(ks.events)
|
||||
case <-time.After(8 * time.Second):
|
||||
slog.Warn("kimiSession: close timed out, abandoning wg.Wait")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func truncate(s string, maxRunes int) string {
|
||||
if utf8.RuneCountInString(s) <= maxRunes {
|
||||
return s
|
||||
}
|
||||
return string([]rune(s)[:maxRunes]) + "..."
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package kimi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNewKimiSession(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, err := newKimiSession(ctx, "kimi", "/tmp", "kimi-k2", "default", "resume-123", nil, 0)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ks)
|
||||
assert.True(t, ks.Alive())
|
||||
assert.Equal(t, "resume-123", ks.CurrentSessionID())
|
||||
|
||||
err = ks.Close()
|
||||
assert.NoError(t, err)
|
||||
assert.False(t, ks.Alive())
|
||||
}
|
||||
|
||||
func TestExtractResumeSessionID(t *testing.T) {
|
||||
cases := []struct {
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{"To resume this session: kimi -r e3690555-60eb-4d50-874b-e3647e9cee5b", "e3690555-60eb-4d50-874b-e3647e9cee5b"},
|
||||
{"To resume this session: kimi --resume abc-def", ""},
|
||||
{"To resume this session: no-id-here", ""},
|
||||
{"random text", ""},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
assert.Equal(t, c.expected, extractResumeSessionID(c.input), "input: %s", c.input)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleAssistantWithText(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, _ := newKimiSession(ctx, "kimi", "/tmp", "", "default", "", nil, 0)
|
||||
defer ks.Close()
|
||||
|
||||
ks.handleEvent(map[string]any{
|
||||
"role": "assistant",
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "Hello!"},
|
||||
},
|
||||
})
|
||||
|
||||
// pendingMsgs should buffer the text
|
||||
assert.Len(t, ks.pendingMsgs, 1)
|
||||
assert.Equal(t, "Hello!", ks.pendingMsgs[0])
|
||||
}
|
||||
|
||||
func TestHandleAssistantWithThink(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, _ := newKimiSession(ctx, "kimi", "/tmp", "", "default", "", nil, 0)
|
||||
defer ks.Close()
|
||||
|
||||
ks.handleEvent(map[string]any{
|
||||
"role": "assistant",
|
||||
"content": []any{
|
||||
map[string]any{"type": "think", "think": "Let me think..."},
|
||||
map[string]any{"type": "text", "text": "Done!"},
|
||||
},
|
||||
})
|
||||
|
||||
events := drainEvents(ks.events, 2)
|
||||
require.Len(t, events, 1)
|
||||
assert.Equal(t, core.EventThinking, events[0].Type)
|
||||
assert.Equal(t, "Let me think...", events[0].Content)
|
||||
assert.Equal(t, "Done!", ks.pendingMsgs[0])
|
||||
}
|
||||
|
||||
func TestHandleAssistantWithToolCalls(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, _ := newKimiSession(ctx, "kimi", "/tmp", "", "default", "", nil, 0)
|
||||
defer ks.Close()
|
||||
|
||||
ks.handleEvent(map[string]any{
|
||||
"role": "assistant",
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "I will run a command"},
|
||||
},
|
||||
"tool_calls": []any{
|
||||
map[string]any{
|
||||
"id": "tool_abc",
|
||||
"function": map[string]any{
|
||||
"name": "Shell",
|
||||
"arguments": `{"command":"echo hello"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
events := drainEvents(ks.events, 3)
|
||||
require.Len(t, events, 2)
|
||||
assert.Equal(t, core.EventThinking, events[0].Type)
|
||||
assert.Equal(t, "I will run a command", events[0].Content)
|
||||
assert.Equal(t, core.EventToolUse, events[1].Type)
|
||||
assert.Equal(t, "Shell", events[1].ToolName)
|
||||
assert.Equal(t, `{"command":"echo hello"}`, events[1].ToolInput)
|
||||
assert.Equal(t, "tool_abc", events[1].RequestID)
|
||||
}
|
||||
|
||||
func TestHandleTool(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, _ := newKimiSession(ctx, "kimi", "/tmp", "", "default", "", nil, 0)
|
||||
defer ks.Close()
|
||||
|
||||
ks.handleEvent(map[string]any{
|
||||
"role": "tool",
|
||||
"tool_call_id": "tool_abc",
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "hello\n"},
|
||||
},
|
||||
})
|
||||
|
||||
events := drainEvents(ks.events, 1)
|
||||
require.Len(t, events, 1)
|
||||
assert.Equal(t, core.EventToolResult, events[0].Type)
|
||||
assert.Equal(t, "tool_abc", events[0].ToolName)
|
||||
assert.Contains(t, events[0].ToolResult, "hello")
|
||||
}
|
||||
|
||||
func TestFlushPendingAsText(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, _ := newKimiSession(ctx, "kimi", "/tmp", "", "default", "", nil, 0)
|
||||
defer ks.Close()
|
||||
|
||||
ks.pendingMsgs = []string{"Hello", " ", "world"}
|
||||
ks.flushPendingAsText()
|
||||
|
||||
events := drainEvents(ks.events, 1)
|
||||
require.Len(t, events, 1)
|
||||
assert.Equal(t, core.EventText, events[0].Type)
|
||||
assert.Equal(t, "Hello world", events[0].Content)
|
||||
assert.Empty(t, ks.pendingMsgs)
|
||||
}
|
||||
|
||||
func TestFlushPendingAsThinking(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ks, _ := newKimiSession(ctx, "kimi", "/tmp", "", "default", "", nil, 0)
|
||||
defer ks.Close()
|
||||
|
||||
ks.pendingMsgs = []string{"Thinking..."}
|
||||
ks.flushPendingAsThinking()
|
||||
|
||||
events := drainEvents(ks.events, 1)
|
||||
require.Len(t, events, 1)
|
||||
assert.Equal(t, core.EventThinking, events[0].Type)
|
||||
assert.Equal(t, "Thinking...", events[0].Content)
|
||||
}
|
||||
|
||||
func TestTruncate(t *testing.T) {
|
||||
assert.Equal(t, "hello", truncate("hello", 10))
|
||||
assert.Equal(t, "hello world", truncate("hello world", 11))
|
||||
assert.Equal(t, "hello worl...", truncate("hello world", 10))
|
||||
}
|
||||
|
||||
func drainEvents(ch <-chan core.Event, max int) []core.Event {
|
||||
var events []core.Event
|
||||
timeout := time.After(500 * time.Millisecond)
|
||||
for i := 0; i < max; i++ {
|
||||
select {
|
||||
case evt := <-ch:
|
||||
events = append(events, evt)
|
||||
case <-timeout:
|
||||
return events
|
||||
}
|
||||
}
|
||||
return events
|
||||
}
|
||||
Reference in New Issue
Block a user