diff --git a/pkg/agent/core.go b/pkg/agent/core.go index fff5fe2..5be659c 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -253,6 +253,22 @@ func (s *agentCore) prompt() string { // Prompt returns the display prompt string for the current session state. func (s *agentCore) Prompt() string { return s.prompt() } +// defaultContextLength is used when the backend cannot report the model's +// actual context window (e.g. CodeWhisperer). 128k tokens is a safe default +// for modern models. +const defaultContextLength = 128000 + +// autoCompactLimit returns the token threshold for auto-compaction. +// Uses 75% of the model's context length, reserving room for output and the +// compaction prompt. +func (s *agentCore) autoCompactLimit(ctx context.Context) int { + ctxLen := s.loopcfg.Backend.ContextLength(ctx) + if ctxLen <= 0 { + ctxLen = defaultContextLength + } + return ctxLen * 3 / 4 +} + func (s *agentCore) saveSession() { if s.session == nil || s.sessionID == "" || s.sessionsDir == "" { return @@ -302,6 +318,14 @@ func (s *agentCore) Submit(ctx context.Context, input string, handler EventHandl s.loopcfg.Output = handler + // Auto-compact before the turn if approaching the context limit. + if s.session != nil { + if limit := s.autoCompactLimit(ctx); limit > 0 && s.session.estimateTokens() >= limit { + emit(s.loopcfg, Event{Role: "info", Content: "auto-compacting context...\n"}) + s.session.compact(ctx, s.loopcfg.Backend) //nolint:errcheck + } + } + s.hooks.Run(HookAgentSpawn) err := run(actCtx, s.loopcfg, s.session) @@ -364,6 +388,22 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve handler(infoEvent(fmt.Sprintf("switched backend to: %s (model: %s)", be.Name(), be.Model()))) return true + case "/models": + models := s.loopcfg.Backend.Models(ctx) + if len(models) == 0 { + handler(infoEvent("no models available")) + return true + } + current := s.loopcfg.Backend.Model() + for _, m := range models { + marker := " " + if m == current { + marker = "* " + } + handler(infoEvent(marker + m)) + } + return true + case "/model": if len(args) == 0 { handler(infoEvent("error: /model requires an argument (e.g., /model qwen3:8b)")) @@ -536,6 +576,7 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve " /agent [name] - show or switch active agent", " /backend - switch backend (ollama, openai)", " /model - switch model", + " /models - list available models", " /queued [pop|clear] - manage queued prompts", " /compact - summarize conversation and compact context", " /context - show context window debug info", diff --git a/pkg/agent/session.go b/pkg/agent/session.go index c91ed61..cefbfc9 100644 --- a/pkg/agent/session.go +++ b/pkg/agent/session.go @@ -243,6 +243,11 @@ func (s *Session) appendUserMessage(content string) { // contextStatsString returns a one-line human-readable context summary. func (s *Session) contextStatsString() string { + return fmt.Sprintf("context: ~%d tokens (%d msgs)", s.estimateTokens(), len(s.messages)) +} + +// estimateTokens returns a rough token count (~4 chars per token). +func (s *Session) estimateTokens() int { chars := 0 for _, m := range s.messages { chars += len(m.Content) @@ -250,7 +255,7 @@ func (s *Session) contextStatsString() string { chars += len(tc.Name) + len(tc.Arguments) } } - return fmt.Sprintf("context: ~%d tokens (%d msgs)", chars/4, len(s.messages)) + return chars / 4 } // contextDebug returns a multi-line breakdown of the history. diff --git a/pkg/backend/anthropic.go b/pkg/backend/anthropic.go index 33321fc..5096148 100644 --- a/pkg/backend/anthropic.go +++ b/pkg/backend/anthropic.go @@ -31,6 +31,16 @@ func (b *AnthropicBackend) DefaultModel() string { return "claude-sonnet-4-5" } func (b *AnthropicBackend) Model() string { return b.model } func (b *AnthropicBackend) SetModel(m string) { b.model = m } +func (b *AnthropicBackend) ContextLength(_ context.Context) int { return 200000 } + +func (b *AnthropicBackend) Models(_ context.Context) []string { + return []string{ + "claude-sonnet-4-5", + "claude-opus-4", + "claude-3-5-haiku-latest", + } +} + // -- wire types -- type anthropicRequest struct { diff --git a/pkg/backend/backend.go b/pkg/backend/backend.go index c083dd1..1694971 100644 --- a/pkg/backend/backend.go +++ b/pkg/backend/backend.go @@ -86,4 +86,9 @@ type Backend interface { Model() string // SetModel changes the active model. SetModel(model string) + // ContextLength returns the context window size in tokens for the active model. + // Returns 0 if unknown. May make an API call; results should be cached by the implementation. + ContextLength(ctx context.Context) int + // Models returns the list of available model IDs from the provider. + Models(ctx context.Context) []string } diff --git a/pkg/backend/codewhisperer.go b/pkg/backend/codewhisperer.go index 237d8cc..e56b1a3 100644 --- a/pkg/backend/codewhisperer.go +++ b/pkg/backend/codewhisperer.go @@ -35,6 +35,7 @@ type CodeWhispererBackend struct { authSource kiroAuthSource authInitErr error httpClient *http.Client + ctxLength int // cached } // NewCodeWhisperer returns a CodeWhisperer backend. See package doc for apiKey semantics. @@ -56,6 +57,59 @@ func (b *CodeWhispererBackend) Name() string { return "kiro" } func (b *CodeWhispererBackend) DefaultModel() string { return "auto" } func (b *CodeWhispererBackend) Model() string { return b.model } func (b *CodeWhispererBackend) SetModel(m string) { b.model = m } +func (b *CodeWhispererBackend) Models(ctx context.Context) []string { + resp := b.fetchModels(ctx) + if resp == nil { + return nil + } + ids := make([]string, len(resp.Models)) + for i, m := range resp.Models { + ids[i] = m.ModelID + } + return ids +} + +func (b *CodeWhispererBackend) fetchModels(ctx context.Context) *kiroListModelsResponse { + if b.authInitErr != nil { + return nil + } + token, err := b.authSource.AccessToken(ctx) + if err != nil { + return nil + } + endpoint, err := b.resolveEndpoint(ctx) + if err != nil { + return nil + } + profileARN, _ := b.authSource.ProfileARN(ctx) + client := newKiroAPIClient(endpoint, token, b.extraHeaders, b.httpClient) + resp, err := client.ListModels(ctx, profileARN) + if err != nil { + return nil + } + return resp +} + +func (b *CodeWhispererBackend) ContextLength(ctx context.Context) int { + if b.ctxLength > 0 { + return b.ctxLength + } + resp := b.fetchModels(ctx) + if resp == nil { + return 0 + } + for _, m := range resp.Models { + if m.TokenLimits != nil && (b.model == "auto" || m.ModelID == b.model) { + b.ctxLength = m.TokenLimits.MaxInputTokens + return b.ctxLength + } + } + if resp.DefaultModel != nil && resp.DefaultModel.TokenLimits != nil { + b.ctxLength = resp.DefaultModel.TokenLimits.MaxInputTokens + return b.ctxLength + } + return 0 +} // ── Backend interface ───────────────────────────────────────────────────────── diff --git a/pkg/backend/codewhisperer_internal.go b/pkg/backend/codewhisperer_internal.go index af6b540..5382700 100644 --- a/pkg/backend/codewhisperer_internal.go +++ b/pkg/backend/codewhisperer_internal.go @@ -608,3 +608,70 @@ func (c *kiroOIDCClient) RefreshToken(ctx context.Context, req kiroRefreshTokenR } return &result, nil } + +// ── ListAvailableModels ─────────────────────────────────────────────────────── + +const ( + kiroListModelsTarget = "AmazonCodeWhispererService.ListAvailableModels" + kiroListModelsUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererruntime/0.1.14474 os/linux lang/rust/1.92.0 md/appVersion-1.27.2 app/AmazonQ-For-CLI" + kiroListModelsAmzUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererruntime/0.1.14474 os/linux lang/rust/1.92.0 m/F,C app/AmazonQ-For-CLI" +) + +type kiroListModelsRequest struct { + Origin string `json:"origin"` + ProfileARN string `json:"profileArn,omitempty"` +} + +type kiroListModelsResponse struct { + DefaultModel *kiroAvailableModel `json:"defaultModel,omitempty"` + Models []kiroAvailableModel `json:"models,omitempty"` +} + +type kiroAvailableModel struct { + ModelID string `json:"modelId,omitempty"` + TokenLimits *kiroTokenLimits `json:"tokenLimits,omitempty"` +} + +type kiroTokenLimits struct { + MaxInputTokens int `json:"maxInputTokens,omitempty"` + MaxOutputTokens int `json:"maxOutputTokens,omitempty"` +} + +func (c *kiroAPIClient) ListModels(ctx context.Context, profileARN string) (*kiroListModelsResponse, error) { + body, err := json.Marshal(kiroListModelsRequest{ + Origin: kiroDefaultOrigin, + ProfileARN: profileARN, + }) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body)) + if err != nil { + return nil, err + } + invocationID, _ := kiroNewUUID() + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("X-Amz-Target", kiroListModelsTarget) + req.Header.Set("Authorization", "Bearer "+c.token) + req.Header.Set("Accept", "*/*") + req.Header.Set("User-Agent", kiroListModelsUserAgent) + req.Header.Set("X-Amz-User-Agent", kiroListModelsAmzUserAgent) + req.Header.Set("Amz-Sdk-Request", kiroAmzSDKRequest) + req.Header.Set("Amz-Sdk-Invocation-Id", invocationID) + for k, v := range c.extraHeaders { + req.Header.Set(k, v) + } + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, kiroDecodeHTTPError(resp) + } + var result kiroListModelsResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, err + } + return &result, nil +} diff --git a/pkg/backend/ollama.go b/pkg/backend/ollama.go index 22dd0ca..d1c02ca 100644 --- a/pkg/backend/ollama.go +++ b/pkg/backend/ollama.go @@ -8,13 +8,16 @@ import ( "fmt" "io" "net/http" + "strings" ) // OllamaBackend speaks the Ollama /api/chat wire format. type OllamaBackend struct { - baseURL string - model string - client *http.Client + baseURL string + model string + client *http.Client + ctxLength int + ctxModel string } func NewOllama(baseURL string) *OllamaBackend { @@ -29,7 +32,70 @@ func NewOllama(baseURL string) *OllamaBackend { func (b *OllamaBackend) Name() string { return "ollama" } func (b *OllamaBackend) DefaultModel() string { return "qwen3.5:9b" } func (b *OllamaBackend) Model() string { return b.model } -func (b *OllamaBackend) SetModel(m string) { b.model = m } +func (b *OllamaBackend) SetModel(m string) { b.model = m; b.ctxLength = 0 } + +func (b *OllamaBackend) Models(ctx context.Context) []string { + req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL+"/api/tags", nil) + if err != nil { + return nil + } + resp, err := b.client.Do(req) + if err != nil || resp.StatusCode != 200 { + if resp != nil { + resp.Body.Close() + } + return nil + } + defer resp.Body.Close() + var result struct { + Models []struct { + Name string `json:"name"` + } `json:"models"` + } + if json.NewDecoder(resp.Body).Decode(&result) != nil { + return nil + } + names := make([]string, len(result.Models)) + for i, m := range result.Models { + names[i] = m.Name + } + return names +} + +func (b *OllamaBackend) ContextLength(ctx context.Context) int { + if b.ctxLength > 0 && b.ctxModel == b.model { + return b.ctxLength + } + body, _ := json.Marshal(map[string]string{"name": b.model}) + req, err := http.NewRequestWithContext(ctx, "POST", b.baseURL+"/api/show", bytes.NewReader(body)) + if err != nil { + return 0 + } + resp, err := b.client.Do(req) + if err != nil || resp.StatusCode != 200 { + if resp != nil { + resp.Body.Close() + } + return 0 + } + defer resp.Body.Close() + var result struct { + ModelInfo map[string]any `json:"model_info"` + } + if json.NewDecoder(resp.Body).Decode(&result) != nil { + return 0 + } + for k, v := range result.ModelInfo { + if strings.HasSuffix(k, ".context_length") { + if f, ok := v.(float64); ok && f > 0 { + b.ctxLength = int(f) + b.ctxModel = b.model + return b.ctxLength + } + } + } + return 0 +} // -- wire types -- diff --git a/pkg/backend/openai.go b/pkg/backend/openai.go index 41075a8..db3655f 100644 --- a/pkg/backend/openai.go +++ b/pkg/backend/openai.go @@ -22,6 +22,14 @@ type OpenAIBackend struct { model string client *http.Client extraHeaders map[string]string // optional; applied after Authorization + ctxLength int // cached; 0 = not yet fetched + ctxModel string // model the cached value is for + cachedModels []openAIModelInfo // cached model list +} + +type openAIModelInfo struct { + ID string `json:"id"` + ContextLength int `json:"context_length"` } func NewOpenAI(name, baseURL, apiKey string) *OpenAIBackend { @@ -35,7 +43,59 @@ func NewOpenAI(name, baseURL, apiKey string) *OpenAIBackend { func (b *OpenAIBackend) Name() string { return b.name } func (b *OpenAIBackend) Model() string { return b.model } -func (b *OpenAIBackend) SetModel(m string) { b.model = m } +func (b *OpenAIBackend) SetModel(m string) { b.model = m; b.ctxLength = 0 } + +func (b *OpenAIBackend) fetchModels(ctx context.Context) []openAIModelInfo { + if len(b.cachedModels) > 0 { + return b.cachedModels + } + req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL+"/v1/models", nil) + if err != nil { + return nil + } + if b.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+b.apiKey) + } + resp, err := b.client.Do(req) + if err != nil || resp.StatusCode != 200 { + if resp != nil { + resp.Body.Close() + } + return nil + } + defer resp.Body.Close() + var result struct { + Data []openAIModelInfo `json:"data"` + } + if json.NewDecoder(resp.Body).Decode(&result) != nil { + return nil + } + b.cachedModels = result.Data + return b.cachedModels +} + +func (b *OpenAIBackend) ContextLength(ctx context.Context) int { + if b.ctxLength > 0 && b.ctxModel == b.model { + return b.ctxLength + } + for _, m := range b.fetchModels(ctx) { + if m.ID == b.model { + b.ctxLength = m.ContextLength + b.ctxModel = b.model + return b.ctxLength + } + } + return 0 +} + +func (b *OpenAIBackend) Models(ctx context.Context) []string { + models := b.fetchModels(ctx) + ids := make([]string, len(models)) + for i, m := range models { + ids[i] = m.ID + } + return ids +} func (b *OpenAIBackend) DefaultModel() string { switch b.name { case "openrouter":