Repository: GoogleCloudPlatform/kubectl-ai Branch: main Commit: 54037aecc023 Files: 152 Total size: 899.0 KB Directory structure: gitextract_bjyvth3y/ ├── .github/ │ ├── ISSUE_TEMPLATE/ │ │ ├── bug_report.md │ │ └── feature_request.md │ ├── actions/ │ │ └── kind-cluster-setup/ │ │ └── action.yaml │ ├── kubectl-ai.cast │ └── workflows/ │ ├── ci-periodic.yaml │ ├── ci-presubmit.yaml │ ├── k8s-bench-evals.yaml │ └── release.yaml ├── .gitignore ├── .goreleaser.yaml ├── .krew.yaml ├── CONTAINER.md ├── LICENSE ├── README.md ├── cmd/ │ ├── main.go │ ├── mcp.go │ └── mcp_test.go ├── contributing.md ├── dev/ │ ├── ci/ │ │ ├── periodics/ │ │ │ ├── analyze-evals.sh │ │ │ └── run-evals.sh │ │ └── presubmits/ │ │ ├── go-build.sh │ │ ├── go-vet.sh │ │ ├── verify-autogen.sh │ │ ├── verify-format.sh │ │ ├── verify-gomod.sh │ │ └── verify-mocks.sh │ └── tasks/ │ ├── build-images │ ├── demo.md │ ├── deploy-to-gke │ ├── deploy-to-kind │ ├── format.sh │ ├── generate-github-actions.sh │ └── gomod.sh ├── docs/ │ ├── bedrock.md │ ├── gke-deployment.md │ ├── mcp-client.md │ ├── mcp-server.md │ ├── mocking.md │ ├── tool-samples/ │ │ ├── argocd.yaml │ │ ├── gcloud.yaml │ │ ├── gh.yaml │ │ └── kustomize.yaml │ └── tools.md ├── go.mod ├── go.sum ├── go.work ├── go.work.sum ├── gollm/ │ ├── README.md │ ├── anthropic.go │ ├── anthropic_test.go │ ├── azopenai.go │ ├── bedrock.go │ ├── factory.go │ ├── factory_test.go │ ├── gemini.go │ ├── go.mod │ ├── go.sum │ ├── grok.go │ ├── http_journal.go │ ├── interfaces.go │ ├── llamacpp.go │ ├── ollama.go │ ├── openai.go │ ├── openai_response.go │ ├── openai_test.go │ ├── persist.go │ ├── schema.go │ └── shims.go ├── images/ │ └── kubectl-ai/ │ └── Dockerfile ├── install.sh ├── internal/ │ └── mocks/ │ ├── agent_mock.go │ ├── generate.go │ ├── gollm_mock.go │ └── tools_mock.go ├── k8s/ │ ├── all_in_one.yaml │ ├── kubectl-ai-gke.yaml │ ├── kubectl-ai.yaml │ └── sandbox/ │ ├── all-in-one.yaml │ ├── cluster_role.yaml │ ├── cluster_role_binding.yaml │ ├── namespace.yaml │ ├── role.yaml │ ├── role_binding.yaml │ └── service_account.yaml ├── kubectl-utils/ │ ├── README.md │ ├── cmd/ │ │ └── kubectl-expect/ │ │ └── main.go │ ├── go.mod │ ├── go.sum │ └── pkg/ │ ├── kel/ │ │ ├── expression.go │ │ └── info.go │ └── kube/ │ ├── client.go │ └── discovery.go ├── makefile ├── modelserving/ │ ├── .gitignore │ ├── README.md │ ├── dev/ │ │ └── tasks/ │ │ ├── build-images │ │ ├── deploy-to-gke │ │ ├── deploy-to-kind │ │ ├── download-model │ │ └── run-local │ ├── images/ │ │ ├── llamacpp-gemma3-12b-it/ │ │ │ └── Dockerfile │ │ └── llamacpp-server/ │ │ └── Dockerfile │ └── k8s/ │ ├── llm-server-cpu.yaml │ ├── llm-server-rpc.yaml │ ├── llm-server.yaml │ ├── rpc-server-cpu.yaml │ └── rpc-server-cuda.yaml └── pkg/ ├── agent/ │ ├── agent_e2e_test.go │ ├── conversation.go │ ├── conversation_test.go │ ├── manager.go │ ├── mcp_client.go │ └── systemprompt_template_default.txt ├── api/ │ └── models.go ├── journal/ │ ├── context.go │ ├── loader.go │ ├── log.go │ └── recorder.go ├── mcp/ │ ├── README.md │ ├── client.go │ ├── config.go │ ├── constants.go │ ├── default_config.yaml │ ├── http_client.go │ ├── interfaces.go │ ├── manager.go │ ├── stdio_client.go │ └── utils.go ├── sandbox/ │ ├── executor.go │ ├── kubernetes.go │ ├── local.go │ ├── seatbelt_executor.go │ └── seatbelt_executor_others.go ├── sessions/ │ ├── filesystem.go │ ├── manager.go │ ├── memory.go │ └── store.go ├── tools/ │ ├── bash_tool.go │ ├── custom_tool.go │ ├── custom_tool_test.go │ ├── interfaces.go │ ├── kubectl_filter.go │ ├── kubectl_filter_test.go │ ├── kubectl_tool.go │ ├── mcp_tool.go │ ├── streaming.go │ └── tools.go └── ui/ ├── html/ │ ├── htmlui.go │ └── index.html ├── interfaces.go ├── terminal.go └── tui.go ================================================ FILE CONTENTS ================================================ ================================================ FILE: .github/ISSUE_TEMPLATE/bug_report.md ================================================ --- name: 🐛 Bug Report about: Report a reproducible bug or issue title: "[Bug]: " labels: bug --- **Environment (please complete the following):** - OS: [e.g. Ubuntu 22.04] - kubectl-ai version (run `kubectl-ai version`): [e.g. 0.3.0] - LLM provider: [e.g. gemini, openai, grok...] - LLM model: [e.g. gemini-2.5-pro] **Describe the bug** A clear and concise description of what the bug is. **To Reproduce** Steps to reproduce the behavior: 1. Go to '...' 2. Run command '...' 3. See error **Expected behavior** What you expected to happen. **Additional context** Add any other context or logs here. ================================================ FILE: .github/ISSUE_TEMPLATE/feature_request.md ================================================ --- name: 🚀 Feature Request about: Suggest an idea for a new feature or improvement title: "[Feature]: " labels: enhancement --- **Is your feature request related to a problem? Please describe.** A clear description of the problem you're trying to solve. **Describe the solution you'd like** A clear description of what you want to happen. **Describe alternatives you've considered** Any alternative solutions or features you’ve thought of. **Additional context** Add any other context, links, or screenshots here. ================================================ FILE: .github/actions/kind-cluster-setup/action.yaml ================================================ name: Kind Cluster Setup description: "Sets up a Kind Kubernetes cluster and authenticates with GCP" inputs: cluster_name: description: "The name of the Kind cluster" required: false default: "periodic-eval-cluster" runs: using: "composite" steps: - uses: actions/checkout@v4 - name: Create k8s Kind Cluster uses: helm/kind-action@v1.12.0 with: cluster_name: ${{ inputs.cluster_name }} wait: 300s - uses: "google-github-actions/auth@v2" with: project_id: "sunilarora-fp" workload_identity_provider: "projects/512195022720/locations/global/workloadIdentityPools/github/providers/kubectl-ai" ================================================ FILE: .github/kubectl-ai.cast ================================================ {"version": 2, "width": 190, "height": 49, "timestamp": 1744912218, "env": {"SHELL": "/bin/bash", "TERM": "screen-256color"}} [0.017789, "o", "bash-3.2$ "] [8.320293, "o", "."] [8.513894, "o", "/"] [8.769729, "o", "k"] [8.993115, "o", "u"] [9.054592, "o", "b"] [9.209666, "o", "e"] [9.861315, "o", "c"] [10.073644, "o", "t"] [10.307934, "o", "\u0007"] [10.308033, "o", "l-"] [10.853505, "o", "a"] [10.969856, "o", "i"] [11.170596, "o", " "] [14.210904, "o", "\r\n"] [14.351719, "o", "\r\n\u001b[38;5;252m\u001b[0m\u001b[38;5;252m\u001b[0m \u001b[38;5;252mHey there, what can I help you with\u001b[0m\u001b[38;5;252m today?\u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\u001b[38;5;252m \u001b[0m\r\n\r\n\r\n\r\n>>> "] [18.72908, "o", "h"] [18.977149, "o", "o"] [19.130497, "o", "w"] [19.379084, "o", "'"] [19.442264, "o", "s"] [19.535301, "o", " "] [19.720667, "o", "n"] [19.945368, "o", "g"] [20.007533, "o", "i"] [20.078888, "o", "n"] [20.191491, "o", "x"] [20.318561, "o", " "] [20.414191, "o", "a"] [20.509116, "o", "p"] [20.640961, "o", "p"] [21.018349, "o", " "] [21.121216, "o", "d"] [21.224431, "o", "o"] [21.42611, "o", "i"] [21.484012, "o", "n"] [21.554566, "o", "g"] [21.595759, "o", " "] [21.692772, "o", "i"] [21.743452, "o", "n"] [21.831695, "o", " "] [21.953349, "o", "m"] [22.153589, "o", "y"] [22.188041, "o", " "] [22.286241, "o", "c"] [22.390762, "o", "l"] [22.571458, "o", "u"] [22.625392, "o", "s"] [22.834233, "o", "t"] [23.01689, "o", "e"] [23.15806, "o", "r"] [23.247809, "o", " "] [23.442082, "o", "?"] [23.657884, "o", "\r\n"] [32.871756, "o", "\u001b[31mError: reading streaming LLM response: iterateResponseStream: invalid stream chunk: {\r\n \"error\": {\r\n \"code\": 500,\r\n \"message\": \"An internal error has occurred. Please retry or report in https://developers.generativeai.google/guide/troubleshooting\",\r\n \"status\": \"INTERNAL\",\r\n \"details\": [\r\n {\r\n \"@type\": \"type.googleapis.com/google.rpc.DebugInfo\",\r\n \"detail\": \"syntax error in the model generated tool call.\"\r\n }\r\n ]\r\n }\r\n}\r\n\r\n\u001b[0m\r\n\r\n>>> "] [52.693406, "o", "^C"] [52.69373, "o", "Received signal, shutting down... interrupt\r\n"] [52.702021, "o", "bash-3.2$ "] [53.929981, "o", "w"] [54.073289, "o", "h"] [54.105594, "o", "i"] [54.199234, "o", "c"] [54.278376, "o", "h"] [54.391889, "o", " "] [55.888282, "o", "k"] [56.095038, "o", "u"] [56.172749, "o", "e"] [56.18066, "o", "b"] [56.361241, "o", "c"] [56.569981, "o", "t"] [56.6488, "o", "l"] [56.847693, "o", "-"] [57.289273, "o", "\b\u001b[K"] [57.429839, "o", "\b\u001b[K"] [57.585486, "o", "\b\u001b[K"] [57.717702, "o", "\b\u001b[K"] [57.874424, "o", "\b\u001b[K"] [58.002212, "o", "\b\u001b[K"] [60.177642, "o", "\b\u001b[K"] [60.429335, "o", "\b\u001b[K"] [60.461237, "o", "\b\u001b[K"] [60.494139, "o", "\b\u001b[K"] [60.527807, "o", "\b\u001b[K"] [60.561504, "o", "\b\u001b[K"] [60.599378, "o", "\b\u001b[K"] [60.63316, "o", "\b\u001b[K"] [60.661988, "o", "\u0007"] [60.694941, "o", "\u0007"] [60.72875, "o", "\u0007"] [60.766251, "o", "\u0007"] [61.063912, "o", "g"] [61.136886, "o", "o"] [61.25609, "o", " "] [61.420433, "o", "b"] [61.458215, "o", "u"] [61.664819, "o", "i"] [61.858026, "o", "l"] [61.951196, "o", "d"] [62.030982, "o", "\r\n"] [64.101433, "o", "bash-3.2$ "] [65.730822, "o", "exit\r\n"] ================================================ FILE: .github/workflows/ci-periodic.yaml ================================================ # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Generated by dev/tasks/generate-github-actions name: Run evals periodically on: schedule: # Run every 15 minutes - cron: "*/15 * * * *" workflow_dispatch: # This allows you to manually trigger the workflow from the GitHub UI inputs: reason: description: "Reason for manual trigger" required: false default: "Manual run via UI" jobs: run-eval: if: github.repository == 'GoogleCloudPlatform/kubectl-ai' runs-on: ubuntu-latest timeout-minutes: 12 # Add "id-token" with the intended permissions. permissions: contents: "read" id-token: "write" steps: - uses: actions/checkout@v4 - name: Kind Cluster Setup uses: ./.github/actions/kind-cluster-setup with: cluster_name: periodic-eval-cluster continue-on-error: false timeout-minutes: 3 - name: Run an easy eval run: | for attempt in 1 2; do echo "=== Evaluation attempt $attempt/2 ===" if timeout 4m bash -c 'TEST_ARGS="--llm-provider vertexai --models gemini-2.5-pro --concurrency=1 --task-pattern=scale-" ./dev/ci/periodics/run-evals.sh'; then echo "Evaluation completed successfully on attempt $attempt" break else echo "Attempt $attempt failed or timed out" # Cleanup any hanging processes pkill -f k8s-ai-bench || true pkill -f kubectl-ai || true if [ $attempt -eq 2 ]; then echo "❌ Both attempts failed" exit 1 else echo "Waiting 10 seconds before retry..." sleep 10 fi fi done - name: Analyse results run: | ./dev/ci/periodics/analyze-evals.sh cat ${{ github.workspace }}/.build/k8s-ai-bench.md >> ${GITHUB_STEP_SUMMARY} concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} cancel-in-progress: false ================================================ FILE: .github/workflows/ci-presubmit.yaml ================================================ # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Generated by dev/tasks/generate-github-actions name: ci-presubmit on: pull_request: types: [opened, synchronize, reopened] push: branches: ["main"] jobs: go-build: runs-on: ubuntu-latest timeout-minutes: 60 steps: - uses: actions/checkout@v4 - uses: actions/setup-go@v5 with: go-version-file: 'go.mod' - name: "Run dev/ci/presubmits/go-build.sh" run: | ./dev/ci/presubmits/go-build.sh go-vet: runs-on: ubuntu-latest timeout-minutes: 60 steps: - uses: actions/checkout@v4 - uses: actions/setup-go@v5 with: go-version-file: 'go.mod' - name: "Run dev/ci/presubmits/go-vet.sh" run: | ./dev/ci/presubmits/go-vet.sh verify-autogen: runs-on: ubuntu-latest timeout-minutes: 60 steps: - uses: actions/checkout@v4 - uses: actions/setup-go@v5 with: go-version-file: 'go.mod' - name: "Run dev/ci/presubmits/verify-autogen.sh" run: | ./dev/ci/presubmits/verify-autogen.sh verify-format: runs-on: ubuntu-latest timeout-minutes: 60 steps: - uses: actions/checkout@v4 - uses: actions/setup-go@v5 with: go-version-file: 'go.mod' - name: "Run dev/ci/presubmits/verify-format.sh" run: | ./dev/ci/presubmits/verify-format.sh verify-gomod: runs-on: ubuntu-latest timeout-minutes: 60 steps: - uses: actions/checkout@v4 - uses: actions/setup-go@v5 with: go-version-file: 'go.mod' - name: "Run dev/ci/presubmits/verify-gomod.sh" run: | ./dev/ci/presubmits/verify-gomod.sh verify-mocks: runs-on: ubuntu-latest timeout-minutes: 60 steps: - uses: actions/checkout@v4 - uses: actions/setup-go@v5 with: go-version-file: 'go.mod' - name: "Run dev/ci/presubmits/verify-mocks.sh" run: | ./dev/ci/presubmits/verify-mocks.sh concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} cancel-in-progress: true ================================================ FILE: .github/workflows/k8s-bench-evals.yaml ================================================ # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # This workflow allows PR owners who add new eval tasks to manually run tests on their changes using the GitHub Actions UI. #It is intended for self-service validation of new or modified evals before merging. name: On-Demand k8s-ai-bench Eval Test on: workflow_dispatch: inputs: task_pattern: description: "Task name or glob pattern to test (must not be '*' or empty; e.g. scale-my-task, scale-foo)" required: true jobs: run-eval: runs-on: ubuntu-latest timeout-minutes: 30 permissions: contents: "read" id-token: "write" steps: - name: Validate task_pattern input run: | if [[ -z "${{ github.event.inputs.task_pattern }}" || "${{ github.event.inputs.task_pattern }}" == "*" || "${{ github.event.inputs.task_pattern }}" == "all" ]]; then echo "Error: You must provide a specific task name or pattern. Wildcards or empty values are not allowed." exit 1 fi - uses: actions/checkout@v4 - name: Kind Cluster Setup uses: ./.github/actions/kind-cluster-setup with: cluster_name: ${{ github.head_ref || github.ref_name }} - name: Run evals # In the future, more options (provider/model/tool-use-shim) may be user-selectable. # For now, these are fixed for CI safety and consistency. env: TEST_ARGS: >- --llm-provider ${{ github.event.inputs.llm_provider || 'vertexai' }} \ --models ${{ github.event.inputs.model || 'gemini-2.5-pro' }} \ --enable-tool-use-shim=${{ github.event.inputs.enable_tool_use_shim || 'false' }} \ --task-pattern=${{ github.event.inputs.task_pattern || 'scale-' }} run: | ./dev/ci/periodics/run-evals.sh - name: Analyse results run: | ./dev/ci/periodics/analyze-evals.sh cat ${{ github.workspace }}/.build/k8s-ai-bench.md >> ${GITHUB_STEP_SUMMARY} concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} cancel-in-progress: false ================================================ FILE: .github/workflows/release.yaml ================================================ name: Release on: push: tags: - 'v*' permissions: contents: write packages: write jobs: goreleaser: runs-on: ubuntu-latest steps: - name: Checkout uses: actions/checkout@v4 with: fetch-depth: 0 - name: Set up Go uses: actions/setup-go@v5 with: check-latest: true - name: Install mockgen run: go install go.uber.org/mock/mockgen@latest - name: Run GoReleaser uses: goreleaser/goreleaser-action@v6 with: distribution: goreleaser version: latest args: release --clean env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - name: Update new version in krew-index uses: rajatjindal/krew-release-bot@v0.0.46 ================================================ FILE: .gitignore ================================================ # Binary ./kubectl-ai bin/ .build/ # Log files *.log trace.log prompt.log app.log # OS specific files .DS_Store .env # IDE specific files .idea/ .vscode/ *.swp *.swo # Go specific *.exe *.test *.prof *.out .aider* # Added by goreleaser init: dist/ # Ignore generated credentials from google-github-actions/auth gha-creds-*.json # air config .air.toml ================================================ FILE: .goreleaser.yaml ================================================ # This is an example .goreleaser.yml file with some sensible defaults. # Make sure to check the documentation at https://goreleaser.com # The lines below are called `modelines`. See `:help modeline` # Feel free to remove those if you don't want/need to use them. # yaml-language-server: $schema=https://goreleaser.com/static/schema.json # vim: set ts=2 sw=2 tw=0 fo=cnqoj version: 2 before: hooks: # You may remove this if you don't use go modules. - go mod tidy # you may remove this if you don't need go generate - go generate ./... builds: - env: - CGO_ENABLED=0 goos: - linux - windows - darwin main: ./cmd ldflags: - -s -w - -X main.version={{.Version}} - -X main.commit={{.Commit}} - -X main.date={{.Date}} archives: - formats: [tar.gz] # this name template makes the OS and Arch compatible with the results of `uname`. name_template: >- {{ .ProjectName }}_ {{- title .Os }}_ {{- if eq .Arch "amd64" }}x86_64 {{- else if eq .Arch "386" }}i386 {{- else }}{{ .Arch }}{{ end }} {{- if .Arm }}v{{ .Arm }}{{ end }} # use zip for windows archives format_overrides: - goos: windows formats: [zip] changelog: sort: asc filters: exclude: - "^docs:" - "^test:" release: footer: >- --- Released by [GoReleaser](https://github.com/goreleaser/goreleaser). ================================================ FILE: .krew.yaml ================================================ apiVersion: krew.googlecontainertools.github.com/v1alpha2 kind: Plugin metadata: name: ai spec: version: {{ .TagName }} homepage: https://github.com/GoogleCloudPlatform/kubectl-ai shortDescription: AI-powered Kubernetes assistant description: | This plugin provides a natural language interface to carry out kubernetes related tasks. The plugin can plan and execute multiple steps given a high level description of a task. It's important to note that this plugin does not replace kubectl. Instead, it makes kubectl accessible to non-kubernetes users and makes kubectl users more productive because now they don't have to remember all the syntax and commands to perform common tasks. caveats: | This plugin uses AI models (LLM) to plan and execute tasks. It supports multiple LLM providers such as Gemini, Azure-OpenAI, Ollama, llamacpp. You can get the API key for the default provider (Gemini) from https://aistudio.google.com/app/apikey. platforms: - selector: matchLabels: os: linux arch: amd64 {{addURIAndSha "https://github.com/GoogleCloudPlatform/kubectl-ai/releases/download/{{ .TagName }}/kubectl-ai_Linux_x86_64.tar.gz" .TagName }} bin: kubectl-ai - selector: matchLabels: os: linux arch: arm64 {{addURIAndSha "https://github.com/GoogleCloudPlatform/kubectl-ai/releases/download/{{ .TagName }}/kubectl-ai_Linux_arm64.tar.gz" .TagName }} bin: kubectl-ai - selector: matchLabels: os: darwin arch: amd64 {{addURIAndSha "https://github.com/GoogleCloudPlatform/kubectl-ai/releases/download/{{ .TagName }}/kubectl-ai_Darwin_x86_64.tar.gz" .TagName }} bin: kubectl-ai - selector: matchLabels: os: darwin arch: arm64 {{addURIAndSha "https://github.com/GoogleCloudPlatform/kubectl-ai/releases/download/{{ .TagName }}/kubectl-ai_Darwin_arm64.tar.gz" .TagName }} bin: kubectl-ai - selector: matchLabels: os: windows arch: amd64 {{addURIAndSha "https://github.com/GoogleCloudPlatform/kubectl-ai/releases/download/{{ .TagName }}/kubectl-ai_Windows_x86_64.zip" .TagName }} bin: kubectl-ai.exe - selector: matchLabels: os: windows arch: arm64 {{addURIAndSha "https://github.com/GoogleCloudPlatform/kubectl-ai/releases/download/{{ .TagName }}/kubectl-ai_Windows_arm64.zip" .TagName }} bin: kubectl-ai.exe ================================================ FILE: CONTAINER.md ================================================ # Running kubectl-ai in a Docker Container ## 1. Build the Docker Image First, clone the `kubectl-ai` repository and build the Docker image from the source code. ```bash git clone https://github.com/GoogleCloudPlatform/kubectl-ai.git cd kubectl-ai docker build -t kubectl-ai:latest -f images/kubectl-ai/Dockerfile . ``` ## 2. Running against a GKE cluster To access a GKE cluster, `kubectl-ai` needs two configurations from your local machine: **Google Cloud credentials** and a **Kubernetes config file**. ### Create Google Cloud Credentials First, create Application Default Credentials [(ADC)](https://cloud.google.com/docs/authentication/application-default-credentials). `kubectl` uses these credentials to authenticate with your GKE cluster. ```bash gcloud auth application-default login ``` This command saves your credentials into the `~/.config/gcloud` directory. ### Configure `kubectl` Next, generate the `kubeconfig` file. This file tells `kubectl` which cluster to connect to and to use your ADC credentials for authentication. ```bash gcloud container clusters get-credentials --location ``` This updates the configuration file at `~/.kube/config`. ## 3. Running the Container Finally, mount both configuration directories into the `kubectl-ai` container when you run it. This example shows how to run `kubectl-ai` with a web interface, mounting all necessary credentials and providing a Gemini API key. ```bash export GEMINI_API_KEY="your_api_key_here" docker run --rm -it -p 8080:8080 \ -v ~/.kube:/root/.kube \ -v ~/.config/gcloud:/root/.config/gcloud \ -e GEMINI_API_KEY \ kubectl-ai:latest \ --ui-listen-address 0.0.0.0:8080 \ --ui-type web ``` Alternatively with the default terminal ui: ```bash export GEMINI_API_KEY="your_api_key_here" docker run --rm -it \ -v ~/.kube:/root/.kube \ -v ~/.config/gcloud:/root/.config/gcloud \ -e GEMINI_API_KEY \ kubectl-ai:latest ``` ================================================ FILE: LICENSE ================================================ Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ================================================ FILE: README.md ================================================ # kubectl-ai [![Go Report Card](https://goreportcard.com/badge/github.com/GoogleCloudPlatform/kubectl-ai)](https://goreportcard.com/report/github.com/GoogleCloudPlatform/kubectl-ai) ![GitHub License](https://img.shields.io/github/license/GoogleCloudPlatform/kubectl-ai) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/GoogleCloudPlatform/kubectl-ai) [![GitHub stars](https://img.shields.io/github/stars/GoogleCloudPlatform/kubectl-ai.svg)](https://github.com/GoogleCloudPlatform/kubectl-ai/stargazers) `kubectl-ai` acts as an intelligent interface, translating user intent into precise Kubernetes operations, making Kubernetes management more accessible and efficient. ![kubectl-ai demo GIF using: kubectl-ai "how's nginx app doing in my cluster"](./.github/kubectl-ai.gif) ## Table of Contents - [Quick Start](#quick-start) - [Installation](#installation) - [Usage](#usage) - [Configuration](#configuration) - [Tools](#tools) - [Docker Quick Start](#docker-quick-start) - [MCP Client Mode](#mcp-client-mode) - [Extras](#extras) - [MCP Server Mode](#mcp-server-mode) - [Start Contributing](#start-contributing) - [Learning Resources](#learning-resources) ## Quick Start First, ensure that kubectl is installed and configured. ### Installation #### Quick Install (Linux & MacOS only) ```shell curl -sSL https://raw.githubusercontent.com/GoogleCloudPlatform/kubectl-ai/main/install.sh | bash ```
Other Installation Methods #### Manual Installation (Linux, MacOS and Windows) 1. Download the latest release from the [releases page](https://github.com/GoogleCloudPlatform/kubectl-ai/releases/latest) for your target machine. 2. Untar the release, make the binary executable and move it to a directory in your $PATH (as shown below). ```shell tar -zxvf kubectl-ai_Darwin_arm64.tar.gz chmod a+x kubectl-ai sudo mv kubectl-ai /usr/local/bin/ ``` #### Install with Krew (Linux/macOS/Windows) First of all, you need to have krew installed, refer to [krew document](https://krew.sigs.k8s.io/docs/user-guide/setup/install/) for more details Then you can install with krew ```shell kubectl krew install ai ``` Now you can invoke `kubectl-ai` as a kubectl plugin like this: `kubectl ai`. #### Install on NixOS There are multiple ways to install `kubectl-ai` on NixOS. For a permanent installation add the following to your NixOS-Configuration: ```nix environment.systemPackages = with pkgs; [ kubectl-ai ]; ``` For a temporary installation, you can use the following command: ```shell nix-shell -p kubectl-ai ```
### Usage `kubectl-ai` supports AI models from `gemini`, `vertexai`, `azopenai`, `openai`, `grok`, `bedrock` and local LLM providers such as `ollama` and `llama.cpp`. #### Using Gemini (Default) Set your Gemini API key as an environment variable. If you don't have a key, get one from [Google AI Studio](https://aistudio.google.com). ```bash export GEMINI_API_KEY=your_api_key_here kubectl-ai # Use different gemini model kubectl-ai --model gemini-2.5-pro-exp-03-25 # Use 2.5 flash (faster) model kubectl-ai --quiet --model gemini-2.5-flash-preview-04-17 "check logs for nginx app in hello namespace" ```
Use other AI models #### Using AI models running locally (ollama or llama.cpp) You can use `kubectl-ai` with AI models running locally. `kubectl-ai` supports [ollama](https://ollama.com/) and [llama.cpp](https://github.com/ggml-org/llama.cpp) to use the AI models running locally. Additionally, the [`modelserving`](modelserving) directory provides tools and instructions for deploying your own `llama.cpp`-based LLM serving endpoints locally or on a Kubernetes cluster. This allows you to host models like Gemma directly in your environment. An example of using Google's `gemma3` model with `ollama`: ```shell # assuming ollama is already running and you have pulled one of the gemma models # ollama pull gemma3:12b-it-qat # if your ollama server is at remote, use OLLAMA_HOST variable to specify the host # export OLLAMA_HOST=http://192.168.1.3:11434/ # enable-tool-use-shim because models require special prompting to enable tool calling kubectl-ai --llm-provider ollama --model gemma3:12b-it-qat --enable-tool-use-shim # you can use `models` command to discover the locally available models >> models ``` #### Using Grok You can use X.AI's Grok model by setting your X.AI API key: ```bash export GROK_API_KEY=your_xai_api_key_here kubectl-ai --llm-provider=grok --model=grok-3-beta ``` #### Using AWS Bedrock You can use AWS Bedrock Claude models with your AWS credentials: ```bash # Configure AWS credentials using AWS SSO aws sso login --profile your-profile-name # Or use other AWS credential methods (IAM roles, environment variables, etc.) # Use Claude 4 Sonnet (default) kubectl-ai --llm-provider=bedrock --model=us.anthropic.claude-sonnet-4-20250514-v1:0 # Use Claude 3.7 Sonnet kubectl-ai --llm-provider=bedrock --model=us.anthropic.claude-3-7-sonnet-20250219-v1:0 # Override model via environment variable export BEDROCK_MODEL=us.anthropic.claude-sonnet-4-20250514-v1:0 kubectl-ai --llm-provider=bedrock ``` AWS Bedrock uses the standard AWS SDK credential chain, supporting: - AWS SSO profiles - IAM roles (for EC2/ECS/Lambda) - Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) - AWS CLI configuration files #### Using Azure OpenAI You can also use Azure OpenAI deployment by setting your OpenAI API key and specifying the provider: ```bash export AZURE_OPENAI_API_KEY=your_azure_openai_api_key_here export AZURE_OPENAI_ENDPOINT=https://your_azure_openai_endpoint_here kubectl-ai --llm-provider=azopenai --model=your_azure_openai_deployment_name_here # or az login kubectl-ai --llm-provider=openai://your_azure_openai_endpoint_here --model=your_azure_openai_deployment_name_here ``` #### Using OpenAI You can also use OpenAI models by setting your OpenAI API key and specifying the provider: ```bash export OPENAI_API_KEY=your_openai_api_key_here kubectl-ai --llm-provider=openai --model=gpt-4.1 ``` #### Using OpenAI Compatible API For example, you can use aliyun qwen-xxx models as follows. ```bash export OPENAI_API_KEY=your_openai_api_key_here export OPENAI_ENDPOINT=https://dashscope.aliyuncs.com/compatible-mode/v1 kubectl-ai --llm-provider=openai --model=qwen-plus ```
Run interactively: ```shell kubectl-ai ``` The interactive mode allows you to have a chat with `kubectl-ai`, asking multiple questions in sequence while maintaining context from previous interactions. Simply type your queries and press Enter to receive responses. To exit the interactive shell, type `exit` or press Ctrl+C. Or, run with a task as input: ```shell kubectl-ai --quiet "fetch logs for nginx app in hello namespace" ``` Combine it with other unix commands: ```shell kubectl-ai < query.txt # OR echo "list pods in the default namespace" | kubectl-ai ``` You can even combine a positional argument with stdin input. The positional argument will be used as a prefix to the stdin content: ```shell cat error.log | kubectl-ai "explain the error" ``` We also support persistence between runs with an opt-in. This lets you save a session to the local filesystem, and resume it to maintain previous context. It even works between different interfaces! ```shell kubectl-ai --new-session # start a new session kubectl-ai --list-sessions # list all saved sessions kubectl-ai --resume-session 20250807-510872 # resume session 20250807-510872 kubectl-ai --delete-session 20250807-510872 # delete session 20250807-510872 ``` ## Configuration You can also configure `kubectl-ai` using a YAML configuration file at `~/.config/kubectl-ai/config.yaml`: ```shell mkdir -p ~/.config/kubectl-ai/ cat < ~/.config/kubectl-ai/config.yaml model: gemini-2.5-flash-preview-04-17 llmProvider: gemini toolConfigPaths: ~/.config/kubectl-ai/tools.yaml EOF ``` Verify your configuration: ```shell kubectl-ai --quiet model ```
More configuration Options Here's a complete configuration file with all available options and their default values: ```yaml # LLM provider configuration llmProvider: "gemini" # Default LLM provider model: "gemini-2.5-pro-preview-06-05" # Default model skipVerifySSL: false # Skip SSL verification for LLM API calls # Tool and permission settings toolConfigPaths: ["~/.config/kubectl-ai/tools.yaml"] # Custom tools configuration paths skipPermissions: false # Skip confirmation for resource-modifying commands enableToolUseShim: false # Enable tool use shim for certain models # MCP configuration mcpServer: false # Run in MCP server mode mcpClient: false # Enable MCP client mode externalTools: false # Discover external MCP tools (requires mcp-server) # Runtime settings maxIterations: 20 # Maximum iterations for the agent quiet: false # Run in non-interactive mode removeWorkdir: false # Remove temporary working directory after execution # Kubernetes configuration kubeconfig: "~/.kube/config" # Path to kubeconfig file # UI configuration uiType: "terminal" # UI mode: "terminal" or "web" uiListenAddress: "localhost:8888" # Address for HTML UI server # Prompt configuration promptTemplateFilePath: "" # Custom prompt template file extraPromptPaths: [] # Additional prompt template paths # Debug and trace settings tracePath: "/tmp/kubectl-ai-trace.txt" # Path to trace file ```
All these settings can be configured through either: 1. Command line flags (e.g., `--model=gemini-2.5-pro`) 2. Configuration file (`~/.config/kubectl-ai/config.yaml`) 3. Environment variables (e.g., `GEMINI_API_KEY`) Command line flags take precedence over configuration file settings. ## Tools `kubectl-ai` leverages LLMs to suggest and execute Kubernetes operations using a set of powerful tools. It comes with built-in tools like `kubectl` and `bash`. You can also extend its capabilities by defining your own custom tools. By default, `kubectl-ai` looks for your tool configurations in `~/.config/kubectl-ai/tools.yaml`. To specify tools configuration files or directories containing tools configuration files, use: ```sh ./kubectl-ai --custom-tools-config= "your prompt here" ``` For further details on how to configure your own tools, [go here](docs/tools.md). ## Docker Quick Start This project provides a Docker image that gives you a standalone environment for running kubectl-ai, including against a GKE cluster. ### Running the container against GKE #### Step 1: Build the Image Clone the repository and build the image with the following command ```bash git clone https://github.com/GoogleCloudPlatform/kubectl-ai.git cd kubectl-ai docker build -t kubectl-ai:latest -f images/kubectl-ai/Dockerfile . ``` #### Step 2: Connect to Your GKE Cluster Set up application default credentials and connect to your GKE cluster. ```bash gcloud auth application-default login # If in a gcloud shell this is not necessary gcloud container clusters get-credentials --zone ``` #### Step 3: Run the kubectl-ai container Below is a sample command that can be used to launch the container with a locally hosted web-ui. Be sure to replace the placeholder values with your specific Google Cloud project ID and location. Note you do not need to mount the gcloud config directory if you're on a cloudshell machine. ```bash docker run --rm -it -p 8080:8080 -v ~/.kube:/root/.kube -v ~/.config/gcloud:/root/.config/gcloud -e GOOGLE_CLOUD_LOCATION=us-central1 -e GOOGLE_CLOUD_PROJECT=my-gcp-project kubectl-ai:latest --llm-provider vertexai --ui-listen-address 0.0.0.0:8080 --ui-type web ``` For more info about running from the container image see [CONTAINER.md](CONTAINER.md) ## MCP Client Mode > **Note:** MCP Client Mode is available in `kubectl-ai` version v0.0.12 and onwards. `kubectl-ai` can connect to external [MCP](https://modelcontextprotocol.io/examples) Servers to access additional tools in addition to built-in tools. ### Quick Start with MCP Client Enable MCP client mode: ```bash kubectl-ai --mcp-client ``` ### MCP Client Configuration Create or edit `~/.config/kubectl-ai/mcp.yaml` to customize MCP servers: ```yaml servers: # Local MCP server (stdio-based) # sequential-thinking: Advanced reasoning and step-by-step analysis - name: sequential-thinking command: npx args: - -y - "@modelcontextprotocol/server-sequential-thinking" # Remote MCP server (HTTP-based) - name: cloudflare-documentation url: https://docs.mcp.cloudflare.com/mcp # Optional: Remote MCP server with authentication - name: custom-api url: https://api.example.com/mcp auth: type: "bearer" token: "${MCP_TOKEN}" ``` The system automatically: - Converts parameter names (snake_case → camelCase) - Handles type conversion (strings → numbers/booleans when appropriate) - Provides fallback behavior for unknown servers No additional setup required - just use the `--mcp-client` flag and the AI will have access to all configured MCP tools. 📖 **For detailed configuration options, troubleshooting, and advanced features for MCP Client mode, see the [MCP Client Documentation](docs/mcp-client.md).** 📖 **For multi-server orchestration and security automation examples, see the [MCP Client Integration Guide](docs/mcp-client.md).** ## Extras You can use the following special keywords for specific actions: - `model`: Display the currently selected model. - `models`: List all available models. - `tools`: List all available tools. - `version`: Display the `kubectl-ai` version. - `reset`: Clear the conversational context. - `clear`: Clear the terminal screen. - `exit` or `quit`: Terminate the interactive shell (Ctrl+C also works). ### Invoking as kubectl plugin You can also run `kubectl ai`. `kubectl` finds any executable file in your `PATH` whose name begins with `kubectl-` as a [plugin](https://kubernetes.io/docs/tasks/extend-kubectl/kubectl-plugins/). ## MCP Server Mode `kubectl-ai` can act as an MCP server that exposes kubectl tools to other MCP clients (like Claude, Cursor, or VS Code). The server can run in two modes: ### Basic MCP Server (Built-in tools only) Expose only kubectl-ai's native Kubernetes tools: ```bash kubectl-ai --mcp-server ``` ### Enhanced MCP Server (With external tool discovery) Additionally discover and expose tools from other MCP servers as a unified interface: ```bash kubectl-ai --mcp-server --external-tools ``` This creates a powerful **tool aggregation hub** where kubectl-ai acts as both: - **MCP Server**: Exposing kubectl tools to clients - **MCP Client**: Consuming tools from other MCP servers To serve clients over HTTP using the streamable transport, run: ```bash kubectl-ai --mcp-server --mcp-server-mode streamable-http --http-port 9080 ``` This starts an MCP endpoint at `http://localhost:9080/mcp`. The enhanced mode provides AI clients with access to both Kubernetes operations and general-purpose tools (filesystem, web search, databases, etc.) through a single MCP endpoint. 📖 **For detailed configuration, examples, and troubleshooting, see the [MCP Server Documentation](docs/mcp-server.md).** ## Start Contributing We welcome contributions to `kubectl-ai` from the community. Take a look at our [contribution guide](contributing.md) to get started. ## Learning Resources ### Talks and Presentations - [From Natural Language to K8s Operations: The MCP Architecture and Practice of kubectl-ai](https://blog.wu-boy.com/2025/10/from-natural-language-to-k8s-operations-the-mcp-architecture-and-practice-of-kubectl-ai-en) - A comprehensive presentation covering the architecture and practical usage of kubectl-ai with MCP (Model Context Protocol). --- *Note: This is not an officially supported Google product. This project is not eligible for the [Google Open Source Software Vulnerability Rewards Program](https://bughunters.google.com/open-source-security).* ================================================ FILE: cmd/main.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package main import ( "bufio" "bytes" "context" "errors" "flag" "fmt" "io" "log" "os" "os/signal" "path/filepath" "slices" "strings" "syscall" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/agent" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sessions" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/ui" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/ui/html" "github.com/spf13/cobra" "github.com/spf13/pflag" "k8s.io/klog/v2" "sigs.k8s.io/yaml" ) // Using the defaults from goreleaser as per https://goreleaser.com/cookbooks/using-main.version/ var ( version = "dev" commit = "none" date = "unknown" ) func BuildRootCommand(opt *Options) (*cobra.Command, error) { rootCmd := &cobra.Command{ Use: "kubectl-ai", Short: "A CLI tool to interact with Kubernetes using natural language", Long: "kubectl-ai is a command-line tool that allows you to interact with your Kubernetes cluster using natural language queries. It leverages large language models to understand your intent and translate it into kubectl", Args: cobra.MaximumNArgs(1), // Only one positional arg is allowed. RunE: func(cmd *cobra.Command, args []string) error { return RunRootCommand(cmd.Context(), *opt, args) }, } rootCmd.AddCommand(&cobra.Command{ Use: "version", Short: "Print the version number of kubectl-ai", Run: func(cmd *cobra.Command, args []string) { fmt.Printf("version: %s\ncommit: %s\ndate: %s\n", version, commit, date) os.Exit(0) }, }) if err := opt.bindCLIFlags(rootCmd.Flags()); err != nil { return nil, err } return rootCmd, nil } type Options struct { ProviderID string `json:"llmProvider,omitempty"` ModelID string `json:"model,omitempty"` // SkipPermissions is a flag to skip asking for confirmation before executing kubectl commands // that modifies resources in the cluster. SkipPermissions bool `json:"skipPermissions,omitempty"` // EnableToolUseShim is a flag to enable tool use shim. // TODO(droot): figure out a better way to discover if the model supports tool use // and set this automatically. EnableToolUseShim bool `json:"enableToolUseShim,omitempty"` // Quiet flag indicates if the agent should run in non-interactive mode. // It requires a query to be provided as a positional argument. Quiet bool `json:"quiet,omitempty"` MCPServer bool `json:"mcpServer,omitempty"` MCPClient bool `json:"mcpClient,omitempty"` // ExternalTools enables discovery and exposure of external MCP tools (only works with --mcp-server) ExternalTools bool `json:"externalTools,omitempty"` MaxIterations int `json:"maxIterations,omitempty"` // MCPServerMode is the mode of the MCP server. only works with --mcp-server. MCPServerMode string `json:"mcpServerMode,omitempty"` // Set the HTTP endpoint port for the MCP server when using HTTP transports like streamable-http. HTTPPort int `json:"httpPort,omitempty"` // KubeConfigPath is the path to the kubeconfig file. // If not provided, the default kubeconfig path will be used. KubeConfigPath string `json:"kubeConfigPath,omitempty"` PromptTemplateFilePath string `json:"promptTemplateFilePath,omitempty"` ExtraPromptPaths []string `json:"extraPromptPaths,omitempty"` TracePath string `json:"tracePath,omitempty"` RemoveWorkDir bool `json:"removeWorkDir,omitempty"` ToolConfigPaths []string `json:"toolConfigPaths,omitempty"` // UIType is the type of user interface to use. UIType ui.Type `json:"uiType,omitempty"` // UIListenAddress is the address to listen for the web UI. UIListenAddress string `json:"uiListenAddress,omitempty"` // SkipVerifySSL is a flag to skip verifying the SSL certificate of the LLM provider. SkipVerifySSL bool `json:"skipVerifySSL,omitempty"` // Session management options ResumeSession string `json:"resumeSession,omitempty"` NewSession bool `json:"newSession,omitempty"` ListSessions bool `json:"listSessions,omitempty"` DeleteSession string `json:"deleteSession,omitempty"` SessionBackend string `json:"sessionBackend,omitempty"` // ShowToolOutput is a flag to disable truncation of tool output in the terminal UI. ShowToolOutput bool `json:"showToolOutput,omitempty"` // Sandbox enables execution of tools in a sandbox environment. // Supported values: "k8s", "seatbelt". // If empty, tools are executed locally. Sandbox string `json:"sandbox,omitempty"` // SandboxImage is the container image to use for the sandbox SandboxImage string `json:"sandboxImage,omitempty"` } var defaultToolConfigPaths = []string{ filepath.Join("{CONFIG}", "kubectl-ai", "tools.yaml"), filepath.Join("{HOME}", ".config", "kubectl-ai", "tools.yaml"), } var defaultConfigPaths = []string{ filepath.Join("{CONFIG}", "kubectl-ai", "config.yaml"), filepath.Join("{HOME}", ".config", "kubectl-ai", "config.yaml"), } func (o *Options) InitDefaults() { o.ProviderID = "gemini" o.ModelID = "gemini-2.5-pro" // by default, confirm before executing kubectl commands that modify resources in the cluster. o.SkipPermissions = false o.MCPServer = false o.MCPClient = false // by default, external tools are disabled (only works with --mcp-server) o.ExternalTools = false // We now default to our strongest model (gemini-2.5-pro-exp-03-25) which supports tool use natively. // so we don't need shim. o.EnableToolUseShim = false o.Quiet = false o.MCPServer = false o.MaxIterations = 20 o.KubeConfigPath = "" o.PromptTemplateFilePath = "" o.ExtraPromptPaths = []string{} o.TracePath = filepath.Join(os.TempDir(), "kubectl-ai-trace.txt") o.RemoveWorkDir = false o.ToolConfigPaths = defaultToolConfigPaths // Default to terminal UI o.UIType = ui.UITypeTerminal // Default UI listen address for HTML UI o.UIListenAddress = "localhost:8888" // Default to not skipping SSL verification o.SkipVerifySSL = false // Default MCP server mode is stdio o.MCPServerMode = "stdio" // Default port for HTTP endpoint when using streamable-http mode o.HTTPPort = 9080 // Session management options o.ResumeSession = "" o.ListSessions = false o.DeleteSession = "" o.SessionBackend = "memory" // By default, hide tool outputs o.ShowToolOutput = false o.Sandbox = "" o.SandboxImage = "bitnami/kubectl:latest" } func (o *Options) LoadConfiguration(b []byte) error { if err := yaml.Unmarshal(b, &o); err != nil { return fmt.Errorf("parsing configuration: %w", err) } return nil } func (o *Options) LoadConfigurationFile() error { configPaths := defaultConfigPaths for _, configPath := range configPaths { pathWithPlaceholdersExpanded := configPath if strings.Contains(pathWithPlaceholdersExpanded, "{CONFIG}") { configDir, err := os.UserConfigDir() if err != nil { return fmt.Errorf("getting user config directory (for config file path %q): %w", configPath, err) } pathWithPlaceholdersExpanded = strings.ReplaceAll(pathWithPlaceholdersExpanded, "{CONFIG}", configDir) } if strings.Contains(pathWithPlaceholdersExpanded, "{HOME}") { homeDir, err := os.UserHomeDir() if err != nil { return fmt.Errorf("getting user home directory (for config file path %q): %w", configPath, err) } pathWithPlaceholdersExpanded = strings.ReplaceAll(pathWithPlaceholdersExpanded, "{HOME}", homeDir) } configPath = filepath.Clean(pathWithPlaceholdersExpanded) configBytes, err := os.ReadFile(configPath) if err != nil { if os.IsNotExist(err) { // ignore missing config files, they are optional } else { fmt.Fprintf(os.Stderr, "warning: could not load defaults from %q: %v\n", configPath, err) } } else if len(configBytes) > 0 { if err := o.LoadConfiguration(configBytes); err != nil { fmt.Fprintf(os.Stderr, "warning: error loading configuration from %q: %v\n", configPath, err) } } } return nil } func main() { ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer cancel() go func() { <-ctx.Done() // restore default behavior for a second signal signal.Stop(make(chan os.Signal)) cancel() klog.Flush() fmt.Fprintf(os.Stderr, "\nReceived signal, shutting down gracefully... (press Ctrl+C again to force)\n") }() if err := run(ctx); err != nil { // Don't print error if it's a context cancellation if !errors.Is(err, context.Canceled) { fmt.Fprintln(os.Stderr, err) } // Exit with non-zero status code on error, unless it's a graceful shutdown. if errors.Is(err, context.Canceled) { os.Exit(0) } os.Exit(1) } } func run(ctx context.Context) error { // klog setup must happen before Cobra parses any flags // add commandline flags for logging klogFlags := flag.NewFlagSet("klog", flag.ExitOnError) klog.InitFlags(klogFlags) klogFlags.Set("logtostderr", "false") klogFlags.Set("log_file", filepath.Join(os.TempDir(), "kubectl-ai.log")) defer klog.Flush() var opt Options opt.InitDefaults() // load YAML config values if err := opt.LoadConfigurationFile(); err != nil { return fmt.Errorf("failed to load config file: %w", err) } rootCmd, err := BuildRootCommand(&opt) if err != nil { return err } // cobra has to know that we pass pass flags with flag lib, otherwise it creates conflict with flags.parse() method // We add just the klog flags we want, not all the klog flags (there are a lot, most of them are very niche) rootCmd.PersistentFlags().AddGoFlag(klogFlags.Lookup("v")) rootCmd.PersistentFlags().AddGoFlag(klogFlags.Lookup("alsologtostderr")) // do this early, before the third-party code logs anything. redirectStdLogToKlog() if err := rootCmd.ExecuteContext(ctx); err != nil { return err } return nil } func (opt *Options) bindCLIFlags(f *pflag.FlagSet) error { f.IntVar(&opt.MaxIterations, "max-iterations", opt.MaxIterations, "maximum number of iterations agent will try before giving up") f.StringVar(&opt.KubeConfigPath, "kubeconfig", opt.KubeConfigPath, "path to kubeconfig file") f.StringVar(&opt.PromptTemplateFilePath, "prompt-template-file-path", opt.PromptTemplateFilePath, "path to custom prompt template file") f.StringArrayVar(&opt.ExtraPromptPaths, "extra-prompt-paths", opt.ExtraPromptPaths, "extra prompt template paths") f.StringVar(&opt.TracePath, "trace-path", opt.TracePath, "path to the trace file") f.BoolVar(&opt.RemoveWorkDir, "remove-workdir", opt.RemoveWorkDir, "remove the temporary working directory after execution") f.StringVar(&opt.ProviderID, "llm-provider", opt.ProviderID, "language model provider") f.StringVar(&opt.ModelID, "model", opt.ModelID, "language model e.g. gemini-2.0-flash-thinking-exp-01-21, gemini-2.0-flash") f.BoolVar(&opt.SkipPermissions, "skip-permissions", opt.SkipPermissions, "(dangerous) skip asking for confirmation before executing kubectl commands that modify resources") f.BoolVar(&opt.MCPServer, "mcp-server", opt.MCPServer, "run in MCP server mode") f.BoolVar(&opt.ExternalTools, "external-tools", opt.ExternalTools, "in MCP server mode, discover and expose external MCP tools") f.StringArrayVar(&opt.ToolConfigPaths, "custom-tools-config", opt.ToolConfigPaths, "path to custom tools config file or directory") f.BoolVar(&opt.MCPClient, "mcp-client", opt.MCPClient, "enable MCP client mode to connect to external MCP servers") f.StringVar(&opt.MCPServerMode, "mcp-server-mode", opt.MCPServerMode, "mode of the MCP server. Supported values: stdio, streamable-http") f.IntVar(&opt.HTTPPort, "http-port", opt.HTTPPort, "port for the HTTP endpoint in MCP server mode (used with --mcp-server when --mcp-server-mode is streamable-http)") f.BoolVar(&opt.EnableToolUseShim, "enable-tool-use-shim", opt.EnableToolUseShim, "enable tool use shim") f.BoolVar(&opt.Quiet, "quiet", opt.Quiet, "run in non-interactive mode, requires a query to be provided as a positional argument") f.Var(&opt.UIType, "ui-type", "user interface type to use. Supported values: terminal, web, tui.") f.StringVar(&opt.UIListenAddress, "ui-listen-address", opt.UIListenAddress, "address to listen for the HTML UI.") f.BoolVar(&opt.SkipVerifySSL, "skip-verify-ssl", opt.SkipVerifySSL, "skip verifying the SSL certificate of the LLM provider") f.BoolVar(&opt.ShowToolOutput, "show-tool-output", opt.ShowToolOutput, "show tool output in the terminal UI") f.StringVar(&opt.Sandbox, "sandbox", opt.Sandbox, "execute tools in a sandbox environment (k8s, seatbelt)") f.StringVar(&opt.SandboxImage, "sandbox-image", opt.SandboxImage, "container image to use for the sandbox") f.StringVar(&opt.ResumeSession, "resume-session", opt.ResumeSession, "ID of session to resume (use 'latest' for the most recent session)") f.BoolVar(&opt.ListSessions, "list-sessions", opt.ListSessions, "list all available sessions") f.StringVar(&opt.DeleteSession, "delete-session", opt.DeleteSession, "delete a session by ID") f.BoolVar(&opt.NewSession, "new-session", opt.NewSession, "start a new persistent session") f.StringVar(&opt.SessionBackend, "session-backend", opt.SessionBackend, "session backend to use (memory or filesystem)") return nil } func RunRootCommand(ctx context.Context, opt Options, args []string) error { var err error // Automatically upgrade backend to filesystem if session persistence flags are requested explicitly if (opt.NewSession || opt.ResumeSession != "" || opt.ListSessions || opt.DeleteSession != "") && opt.SessionBackend == "memory" { klog.Infof("Upgrading session-backend to 'filesystem' based on provided flags") opt.SessionBackend = "filesystem" } // Validate flag combinations if opt.ExternalTools && !opt.MCPServer { return fmt.Errorf("--external-tools can only be used with --mcp-server") } // resolve kubeconfig path with priority: flag/env > KUBECONFIG > default path if err = resolveKubeConfigPath(&opt); err != nil { return fmt.Errorf("failed to resolve kubeconfig path: %w", err) } if opt.MCPServer { if err = startMCPServer(ctx, opt); err != nil { return fmt.Errorf("failed to start MCP server: %w", err) } return nil // MCP server mode blocks, so we return here } if opt.ListSessions { return handleListSessions(opt) } if opt.DeleteSession != "" { return handleDeleteSession(opt) } if err := handleCustomTools(opt.ToolConfigPaths); err != nil { return fmt.Errorf("failed to process custom tools: %w", err) } // After reading stdin, it is consumed var hasInputData bool hasInputData, err = hasStdInData() if err != nil { return fmt.Errorf("failed to check if stdin has data: %w", err) } // Handles positional args or stdin var queryFromCmd string queryFromCmd, err = resolveQueryInput(hasInputData, args) if err != nil { return fmt.Errorf("failed to resolve query input %w", err) } klog.Info("Application started", "pid", os.Getpid()) var recorder journal.Recorder if opt.TracePath != "" { var fileRecorder journal.Recorder fileRecorder, err = journal.NewFileRecorder(opt.TracePath) if err != nil { return fmt.Errorf("creating trace recorder: %w", err) } defer fileRecorder.Close() recorder = fileRecorder } else { // Ensure we always have a recorder, to avoid nil checks recorder = &journal.LogRecorder{} defer recorder.Close() } // Initialize session management var session *api.Session var sessionManager *sessions.SessionManager sessionManager, err = sessions.NewSessionManager(opt.SessionBackend) if err != nil { return fmt.Errorf("failed to create session manager: %w", err) } // Build agentFactory for new agents agentFactory := func(ctx context.Context) (*agent.Agent, error) { var client gollm.Client var err error if opt.SkipVerifySSL { client, err = gollm.NewClient(ctx, opt.ProviderID, gollm.WithSkipVerifySSL()) } else { client, err = gollm.NewClient(ctx, opt.ProviderID) } if err != nil { return nil, fmt.Errorf("creating llm client: %w", err) } return &agent.Agent{ Model: opt.ModelID, Provider: opt.ProviderID, Kubeconfig: opt.KubeConfigPath, LLM: client, MaxIterations: opt.MaxIterations, PromptTemplateFile: opt.PromptTemplateFilePath, ExtraPromptPaths: opt.ExtraPromptPaths, Tools: tools.Default(), Recorder: recorder, RemoveWorkDir: opt.RemoveWorkDir, SkipPermissions: opt.SkipPermissions, EnableToolUseShim: opt.EnableToolUseShim, MCPClientEnabled: opt.MCPClient, Sandbox: opt.Sandbox, SandboxImage: opt.SandboxImage, SessionBackend: opt.SessionBackend, RunOnce: opt.Quiet, InitialQuery: queryFromCmd, }, nil } agentManager := agent.NewAgentManager(agentFactory, sessionManager) // Register cleanup for all sessions and agents defer agentManager.Close() if opt.ResumeSession != "" { if opt.ResumeSession == "latest" { session, err = sessionManager.GetLatestSession() if err != nil { return fmt.Errorf("failed to get latest session: %w", err) } if session == nil { // No latest session found, create a new one klog.Info("No previous session found to resume. Creating new session.") } } else { session, err = sessionManager.FindSessionByID(opt.ResumeSession) if err != nil { return fmt.Errorf("session %s not found: %w", opt.ResumeSession, err) } } } var defaultAgent *agent.Agent // If no session loaded (or resume failed/not requested), create a new one if session == nil { meta := sessions.Metadata{ ModelID: opt.ModelID, ProviderID: opt.ProviderID, } session, err = sessionManager.NewSession(meta) if err != nil { return fmt.Errorf("failed to create a new session: %w", err) } defaultAgent, err = agentManager.GetAgent(ctx, session.ID) if err != nil { return fmt.Errorf("failed to get agent for new session: %w", err) } klog.Infof("Created new session: %s\n", session.ID) } else { // Update last accessed for resumed session if err := sessionManager.UpdateLastAccessed(session); err != nil { klog.Warningf("Failed to update session last accessed time: %v", err) } klog.Infof("Resuming session: %s\n", session.ID) defaultAgent, err = agentManager.GetAgent(ctx, session.ID) if err != nil { return fmt.Errorf("failed to get agent for session: %w", err) } } var userInterface ui.UI switch opt.UIType { case ui.UITypeTerminal: // since stdin is already consumed, we use TTY for taking input from user useTTYForInput := hasInputData userInterface, err = ui.NewTerminalUI(defaultAgent, useTTYForInput, opt.ShowToolOutput, recorder) if err != nil { return fmt.Errorf("creating terminal UI: %w", err) } case ui.UITypeWeb: userInterface, err = html.NewHTMLUserInterface(agentManager, sessionManager, opt.ModelID, opt.ProviderID, opt.UIListenAddress, recorder) if err != nil { return fmt.Errorf("creating web UI: %w", err) } case ui.UITypeTUI: userInterface = ui.NewTUI(defaultAgent) default: return fmt.Errorf("ui-type mode %q is not known", opt.UIType) } err = userInterface.Run(ctx) if err != nil && !errors.Is(err, context.Canceled) { return fmt.Errorf("running UI: %w", err) } return nil } func handleCustomTools(toolConfigPaths []string) error { // resolve tool config paths, and then load and register custom tools from config files and dirs for _, path := range toolConfigPaths { pathWithPlaceholdersExpanded := path if strings.Contains(pathWithPlaceholdersExpanded, "{CONFIG}") { configDir, err := os.UserConfigDir() if err != nil { klog.Warningf("Failed to get user config directory for tools path %q: %v", path, err) continue } pathWithPlaceholdersExpanded = strings.ReplaceAll(pathWithPlaceholdersExpanded, "{CONFIG}", configDir) } if strings.Contains(pathWithPlaceholdersExpanded, "{HOME}") { homeDir, err := os.UserHomeDir() if err != nil { klog.Warningf("Failed to get user home directory for tools path %q: %v", path, err) continue } pathWithPlaceholdersExpanded = strings.ReplaceAll(pathWithPlaceholdersExpanded, "{HOME}", homeDir) } cleanedPath := filepath.Clean(pathWithPlaceholdersExpanded) klog.Infof("Attempting to load custom tools from processed path: %q (original value from config: %q)", cleanedPath, path) if err := tools.LoadAndRegisterCustomTools(cleanedPath); err != nil { if errors.Is(err, os.ErrNotExist) && !slices.Contains(defaultToolConfigPaths, path) { // user specified a directory that does not exist, we must error out return fmt.Errorf("custom tools directory not found (original value: %q, processed path: %q)", path, cleanedPath) } else { klog.Warningf("Failed to load or register custom tools (original value: %q, processed path: %q): %v", path, cleanedPath, err) } } } return nil } // Redirect standard log output to our custom klog writer // This is primarily to suppress warning messages from // genai library https://github.com/googleapis/go-genai/blob/6ac4afc0168762dc3b7a4d940fc463cc1854f366/types.go#L1633 func redirectStdLogToKlog() { log.SetOutput(klogWriter{}) // Disable standard log's prefixes (date, time, file info) // because klog will add its own more detailed prefix. log.SetFlags(0) } // Define a custom writer that forwards messages to klog.Warning type klogWriter struct{} func (writer klogWriter) Write(data []byte) (n int, err error) { // We trim the trailing newline because klog adds its own. message := string(bytes.TrimSuffix(data, []byte("\n"))) klog.Warning(message) return len(data), nil } func hasStdInData() (bool, error) { hasData := false stat, err := os.Stdin.Stat() if err != nil { return hasData, fmt.Errorf("checking stdin: %w", err) } hasData = (stat.Mode() & os.ModeCharDevice) == 0 return hasData, nil } // resolveQueryInput determines the query input from positional args and/or stdin. // It supports: // - 1 positional arg only -> kubectl-ai "get pods" // - stdin only -> echo "get pods" | kubectl-ai // - 1 positional arg + stdin (combined) -> kubectl-ai get <<< "pods" or kubectl-ai "get" <<< "pods" // As default no positional arg nor stdin func resolveQueryInput(hasStdInData bool, args []string) (string, error) { switch { case len(args) == 1 && !hasStdInData: // Use argument directly return args[0], nil case len(args) == 1 && hasStdInData: // Combine arg + stdin var b strings.Builder b.WriteString(args[0]) b.WriteString("\n") scanner := bufio.NewScanner(os.Stdin) for scanner.Scan() { b.WriteString(scanner.Text()) b.WriteString("\n") } if err := scanner.Err(); err != nil { return "", fmt.Errorf("reading stdin: %w", err) } query := strings.TrimSpace(b.String()) if query == "" { return "", fmt.Errorf("no query provided from stdin") } return query, nil case len(args) == 0 && hasStdInData: // Read stdin only b, err := io.ReadAll(os.Stdin) if err != nil { return "", fmt.Errorf("reading stdin: %w", err) } query := strings.TrimSpace(string(b)) if query == "" { return "", fmt.Errorf("no query provided from stdin") } return query, nil default: // Case: No input at all — return empty string, no error return "", nil } } func resolveKubeConfigPath(opt *Options) error { switch { case opt.KubeConfigPath != "": // Already set from flag or viper env case os.Getenv("KUBECONFIG") != "": opt.KubeConfigPath = os.Getenv("KUBECONFIG") default: home, err := os.UserHomeDir() if err != nil { return fmt.Errorf("failed to get user home directory: %w", err) } defaultPath := filepath.Join(home, ".kube", "config") // Only use the default path if it exists if _, err := os.Stat(defaultPath); err == nil { opt.KubeConfigPath = defaultPath } } // We resolve the kubeconfig path to an absolute path, so we can run kubectl from any working directory. if opt.KubeConfigPath != "" { p, err := filepath.Abs(opt.KubeConfigPath) if err != nil { return fmt.Errorf("failed to get absolute path for kubeconfig file %q: %w", opt.KubeConfigPath, err) } opt.KubeConfigPath = p } return nil } func startMCPServer(ctx context.Context, opt Options) error { workDir := filepath.Join(os.TempDir(), "kubectl-ai-mcp") if err := os.MkdirAll(workDir, 0o755); err != nil { return fmt.Errorf("error creating work directory: %w", err) } mcpServer, err := newKubectlMCPServer(ctx, opt.KubeConfigPath, tools.Default(), workDir, opt.ExternalTools, opt.MCPServerMode, opt.HTTPPort) if err != nil { return fmt.Errorf("creating mcp server: %w", err) } return mcpServer.Serve(ctx) } // handleListSessions lists all available sessions with their metadata. func handleListSessions(opt Options) error { manager, err := sessions.NewSessionManager(opt.SessionBackend) if err != nil { return fmt.Errorf("failed to create session manager: %w", err) } sessionList, err := manager.ListSessions() if err != nil { return fmt.Errorf("failed to list sessions: %w", err) } if len(sessionList) == 0 { fmt.Println("No sessions found.") return nil } fmt.Println("Available sessions:") fmt.Println("ID\t\tCreated\t\t\tLast Accessed\t\tModel\t\tProvider") fmt.Println("--\t\t-------\t\t\t-------------\t\t-----\t\t--------") for _, session := range sessionList { fmt.Printf("%s\t%s\t%s\t%s\t%s\n", session.ID, session.CreatedAt.Format("2006-01-02 15:04:05"), session.LastModified.Format("2006-01-02 15:04:05"), session.ModelID, session.ProviderID) } return nil } // handleDeleteSession deletes a session by ID. func handleDeleteSession(opt Options) error { manager, err := sessions.NewSessionManager(opt.SessionBackend) if err != nil { return fmt.Errorf("failed to create session manager: %w", err) } session, err := manager.FindSessionByID(opt.DeleteSession) if err != nil { return fmt.Errorf("session %s not found: %w", opt.DeleteSession, err) } fmt.Printf("Deleting session %s:\n", opt.DeleteSession) fmt.Printf(" Model: %s\n", session.ModelID) fmt.Printf(" Provider: %s\n", session.ProviderID) fmt.Printf(" Created: %s\n", session.CreatedAt.Format("2006-01-02 15:04:05")) fmt.Print("Are you sure you want to delete this session? (y/N): ") var response string fmt.Scanln(&response) if response != "y" && response != "Y" { fmt.Println("Deletion cancelled.") return nil } if err := manager.DeleteSession(opt.DeleteSession); err != nil { return fmt.Errorf("failed to delete session: %w", err) } fmt.Printf("Session %s deleted successfully.\n", opt.DeleteSession) return nil } ================================================ FILE: cmd/mcp.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package main import ( "context" "fmt" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/mcp" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" mcpgo "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" "k8s.io/klog/v2" ) type kubectlMCPServer struct { kubectlConfig string server *server.MCPServer tools tools.Tools workDir string mcpManager *mcp.Manager // Add MCP manager for external tool calls mcpServerMode string // Server mode (e.g., "streamable-http", "stdio") httpPort int // Port for HTTP-based server modes } func newKubectlMCPServer(ctx context.Context, kubectlConfig string, tools tools.Tools, workDir string, exposeExternalTools bool, serverMode string, httpPort int) (*kubectlMCPServer, error) { s := &kubectlMCPServer{ kubectlConfig: kubectlConfig, workDir: workDir, server: server.NewMCPServer( "kubectl-ai", "0.0.1", server.WithToolCapabilities(true), ), tools: tools, mcpServerMode: serverMode, httpPort: httpPort, } // Add built-in tools for _, tool := range s.tools.AllTools() { toolDefn := tool.FunctionDefinition() toolInputSchema, err := toolDefn.Parameters.ToRawSchema() if err != nil { return nil, fmt.Errorf("converting tool schema to json.RawMessage: %w", err) } s.server.AddTool(mcpgo.NewToolWithRawSchema( toolDefn.Name, toolDefn.Description, toolInputSchema, ), s.handleToolCall) } // Only discover external MCP tools if explicitly enabled if exposeExternalTools { // Initialize MCP manager to get client tools manager, err := mcp.InitializeManager() if err != nil { klog.Warningf("Failed to initialize MCP manager: %v", err) return s, nil // Return server with just built-in tools } // Store the manager for later use in tool calls s.mcpManager = manager // Connect to MCP servers and get their tools if err := manager.DiscoverAndConnectServers(ctx); err != nil { klog.Warningf("Failed to connect to MCP servers: %v", err) return s, nil // Return server with just built-in tools } // Get tools from all connected MCP servers serverTools, err := manager.ListAvailableTools(ctx) if err != nil { klog.Warningf("Failed to list tools from MCP servers: %v", err) return s, nil // Return server with just built-in tools } // Add tools from MCP servers totalToolsRegistered := 0 for serverName, tools := range serverTools { klog.V(2).Infof("Processing tools from MCP server %s: %d tools found", serverName, len(tools)) for _, tool := range tools { // Create unique tool name to avoid conflicts with built-in tools or from other servers uniqueToolName := fmt.Sprintf("%s_%s", serverName, tool.Name) // Use the actual tool schema instead of creating a generic wrapper var schema *gollm.FunctionDefinition if tool.InputSchema != nil { // Use the real schema from the external tool schema = &gollm.FunctionDefinition{ Name: uniqueToolName, Description: fmt.Sprintf("%s (from %s)", tool.Description, serverName), Parameters: tool.InputSchema, } } else { // Fallback to generic schema if no schema provided klog.V(2).Infof("External tool %s from %s has no schema, using generic wrapper", tool.Name, serverName) schema = &gollm.FunctionDefinition{ Name: uniqueToolName, Description: fmt.Sprintf("%s (from %s)", tool.Description, serverName), Parameters: &gollm.Schema{ Type: gollm.TypeObject, Properties: map[string]*gollm.Schema{ "args": { Type: gollm.TypeObject, Description: "Tool arguments", }, }, }, } } toolInputSchema, err := schema.Parameters.ToRawSchema() if err != nil { klog.Errorf("Failed to convert tool schema for %s from %s: %v - skipping tool", tool.Name, serverName, err) continue } // Add the tool to the server s.server.AddTool(mcpgo.NewToolWithRawSchema( uniqueToolName, schema.Description, toolInputSchema, ), s.handleToolCall) totalToolsRegistered++ klog.V(3).Infof("Registered tool: %s from server %s", uniqueToolName, serverName) } } klog.Infof("MCP server initialized with external tool discovery enabled - registered %d tools from %d servers", totalToolsRegistered, len(serverTools)) } else { klog.Infof("MCP server initialized with external tool discovery disabled") } return s, nil } func (s *kubectlMCPServer) Serve(ctx context.Context) error { // Ensure proper cleanup of MCP manager on shutdown if s.mcpManager != nil { defer func() { if err := s.mcpManager.Close(); err != nil { klog.Warningf("Failed to close MCP manager: %v", err) } }() } klog.Info("Starting kubectl-ai MCP server") switch s.mcpServerMode { case "streamable-http": // Start the server in streamable HTTP mode klog.Infof("Starting MCP server in streamable HTTP mode on port %d", s.httpPort) httpServer := server.NewStreamableHTTPServer(s.server) endpoint := fmt.Sprintf(":%d", s.httpPort) klog.Infof("Listening for streamable HTTP connections on port %d", s.httpPort) return httpServer.Start(endpoint) default: return server.ServeStdio(s.server) } } func (s *kubectlMCPServer) handleToolCall(ctx context.Context, request mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { toolName := request.Params.Name // First, try to find the tool in our built-in tools collection builtinTool := s.tools.Lookup(toolName) if builtinTool != nil { return s.handleBuiltinToolCall(ctx, request, builtinTool) } // If not a built-in tool, try to handle as external MCP tool if s.mcpManager != nil { return s.handleExternalMCPToolCall(ctx, request) } // Tool not found return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("tool %q not found", toolName), }, }, }, nil } // handleBuiltinToolCall handles calls to built-in kubectl-ai tools func (s *kubectlMCPServer) handleBuiltinToolCall(ctx context.Context, request mcpgo.CallToolRequest, tool tools.Tool) (*mcpgo.CallToolResult, error) { // Set up context for built-in tools ctx = context.WithValue(ctx, tools.KubeconfigKey, s.kubectlConfig) ctx = context.WithValue(ctx, tools.WorkDirKey, s.workDir) // Convert arguments to the expected type args, ok := request.Params.Arguments.(map[string]any) if !ok { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("invalid arguments type: expected map[string]any, got %T", request.Params.Arguments), }, }, }, nil } // Execute the built-in tool result, err := tool.Run(ctx, args) if err != nil { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: err.Error(), }, }, }, nil } // Convert result to string var resultStr string switch v := result.(type) { case string: resultStr = v default: resultStr = fmt.Sprintf("%v", v) } return &mcpgo.CallToolResult{ Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: resultStr, }, }, }, nil } // handleExternalMCPToolCall handles calls to external MCP tools func (s *kubectlMCPServer) handleExternalMCPToolCall(ctx context.Context, request mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { toolName := request.Params.Name // Find which server provides this tool serverTools, err := s.mcpManager.ListAvailableTools(ctx) if err != nil { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("failed to list available tools: %v", err), }, }, }, nil } var targetServerName string var originalToolName string // Look for the tool by checking both original name and server-prefixed name for serverName, tools := range serverTools { for _, tool := range tools { uniqueToolName := fmt.Sprintf("%s_%s", serverName, tool.Name) if uniqueToolName == toolName { targetServerName = serverName originalToolName = tool.Name // Use the original tool name for the MCP call break } } if targetServerName != "" { break } } if targetServerName == "" { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("external MCP tool %q not found", toolName), }, }, }, nil } // Get the client for the target server client, exists := s.mcpManager.GetClient(targetServerName) if !exists { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("MCP client for server %q not found", targetServerName), }, }, }, nil } // Extract arguments - handle the args wrapper for external tools and empty/nil input var toolArgs map[string]any if request.Params.Arguments == nil { // Handle nil arguments as empty map toolArgs = make(map[string]any) } else if args, ok := request.Params.Arguments.(map[string]any); ok { if argsValue, hasArgs := args["args"]; hasArgs { if argsMap, ok := argsValue.(map[string]any); ok { toolArgs = argsMap } else { toolArgs = args // Fallback to using args directly } } else { toolArgs = args // Use arguments directly if no "args" wrapper } } else { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("invalid arguments type: expected map[string]any, got %T", request.Params.Arguments), }, }, }, nil } // Call the external MCP tool using the original tool name result, err := client.CallTool(ctx, originalToolName, toolArgs) if err != nil { return &mcpgo.CallToolResult{ IsError: true, Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: fmt.Sprintf("error calling external MCP tool %q on server %q: %v", originalToolName, targetServerName, err), }, }, }, nil } // Return successful result return &mcpgo.CallToolResult{ Content: []mcpgo.Content{ mcpgo.TextContent{ Type: "text", Text: result, }, }, }, nil } ================================================ FILE: cmd/mcp_test.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package main import ( "context" "fmt" "net" "testing" "time" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/mcp" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" ) func TestKubectlMCPServerHTTPClientIntegration(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() toolset := tools.Tools{} toolset.Init() toolset.RegisterTool(&stubTool{}) port := getFreePort(t) workDir := t.TempDir() server, err := newKubectlMCPServer(ctx, "", toolset, workDir, false, "streamable-http", port) if err != nil { t.Fatalf("failed to create MCP server: %v", err) } serverErr := make(chan error, 1) go func() { serverErr <- server.Serve(ctx) }() waitForHTTPServer(t, port) select { case err := <-serverErr: if err != nil { t.Fatalf("server exited early: %v", err) } default: } clientConfig := mcp.ClientConfig{ Name: "test-client", URL: fmt.Sprintf("http://127.0.0.1:%d/mcp", port), UseStreaming: true, Timeout: 5, } client := mcp.NewClient(clientConfig) connectCtx, connectCancel := context.WithTimeout(ctx, 5*time.Second) defer connectCancel() t.Log("connecting client") if err := client.Connect(connectCtx); err != nil { t.Fatalf("failed to connect client to MCP server: %v", err) } defer func() { if err := client.Close(); err != nil { t.Errorf("failed to close MCP client: %v", err) } }() toolsCtx, toolsCancel := context.WithTimeout(ctx, 5*time.Second) defer toolsCancel() t.Log("listing tools") availableTools, err := client.ListTools(toolsCtx) if err != nil { t.Fatalf("failed to list tools from MCP server: %v", err) } t.Logf("retrieved %d tool(s)", len(availableTools)) if !toolExists("stub", availableTools) { t.Fatalf("expected to find stub tool, got %v", availableTools) } cancel() select { case <-serverErr: case <-time.After(500 * time.Millisecond): } } func waitForHTTPServer(t *testing.T, port int) { t.Helper() deadline := time.Now().Add(5 * time.Second) address := fmt.Sprintf("127.0.0.1:%d", port) for { conn, err := net.DialTimeout("tcp", address, 100*time.Millisecond) if err == nil { conn.Close() return } if time.Now().After(deadline) { t.Fatalf("server did not start listening on %s: %v", address, err) } time.Sleep(50 * time.Millisecond) } } func getFreePort(t *testing.T) int { t.Helper() l, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("failed to acquire free port: %v", err) } defer l.Close() return l.Addr().(*net.TCPAddr).Port } func toolExists(name string, tools []mcp.Tool) bool { for _, tool := range tools { if tool.Name == name { return true } } return false } type stubTool struct{} func (stubTool) Name() string { return "stub" } func (stubTool) Description() string { return "stub tool" } func (stubTool) FunctionDefinition() *gollm.FunctionDefinition { return &gollm.FunctionDefinition{ Name: "stub", Description: "stub tool", Parameters: &gollm.Schema{ Type: gollm.TypeObject, }, } } func (stubTool) Run(context.Context, map[string]any) (any, error) { return "ok", nil } func (stubTool) IsInteractive(map[string]any) (bool, error) { return false, nil } func (stubTool) CheckModifiesResource(map[string]any) string { return "no" } ================================================ FILE: contributing.md ================================================ # How to Contribute We would love to accept your patches and contributions to this project. ## Before you begin ### Sign our Contributor License Agreement Contributions to this project must be accompanied by a [Contributor License Agreement](https://cla.developers.google.com/about) (CLA). You (or your employer) retain the copyright to your contribution; this simply gives us permission to use and redistribute your contributions as part of the project. If you or your current employer have already signed the Google CLA (even if it was for a different project), you probably don't need to do it again. Visit to see your current agreements or to sign a new one. ### Review our Community Guidelines This project follows [Google's Open Source Community Guidelines](https://opensource.google/conduct/). ## Contribution process ### Code Reviews All submissions, including submissions by project members, require review. We use [GitHub pull requests](https://docs.github.com/articles/about-pull-requests) for this purpose. ## Understand the repo An AI-generated overview of the system architecture for this repository is available [here](https://deepwiki.com/GoogleCloudPlatform/kubectl-ai/). This can provide an interactive way to explore the codebase. Quick notes about the various directories: - Source code for `kubectl-ai` CLI lives under `cmd/` and `pkg/` directories. - gollm directory is an independent Go module that implements LLM clients for different LLM providers. - `modelserving` directory contains utilities and configuration to build and run open source AI models locally or in a kubernetes cluster. - `kubectl-utils` is an independent Go package/binary to help with the benchmarks tasks that evaluates various conditions involving properties of kubernetes resources. - User guides/design docs/proposals live under `docs` directory. - `dev` directory scripts for project related tasks (adhoc/CI). ================================================ FILE: dev/ci/periodics/analyze-evals.sh ================================================ #!/bin/bash set -o errexit set -o nounset set -o pipefail set -x REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} if [[ -z "${OUTPUT_DIR:-}" ]]; then OUTPUT_DIR="${REPO_ROOT}/.build/k8s-ai-bench" mkdir -p "${OUTPUT_DIR}" fi BINDIR="${REPO_ROOT}/.build/bin" mkdir -p "${BINDIR}" K8S_AI_BENCH_SRC="${REPO_ROOT}/.build/k8s-ai-bench-src" rm -rf "${K8S_AI_BENCH_SRC}" git clone https://github.com/gke-labs/k8s-ai-bench "${K8S_AI_BENCH_SRC}" cd "${K8S_AI_BENCH_SRC}" GOWORK=off go build -o "${BINDIR}/k8s-ai-bench" . cd "${REPO_ROOT}" # Pass --show-failures flag to the analyze command if it's set ANALYZE_ARGS="" if [[ "$*" == *"--show-failures"* ]]; then ANALYZE_ARGS="--show-failures" fi "${BINDIR}/k8s-ai-bench" analyze --input-dir "${OUTPUT_DIR}" ${TEST_ARGS:-} -results-filepath ${REPO_ROOT}/.build/k8s-ai-bench.md --output-format markdown ${ANALYZE_ARGS} "${BINDIR}/k8s-ai-bench" analyze --input-dir "${OUTPUT_DIR}" ${TEST_ARGS:-} -results-filepath ${REPO_ROOT}/.build/k8s-ai-bench.json --output-format json ================================================ FILE: dev/ci/periodics/run-evals.sh ================================================ #!/bin/bash set -o errexit set -o nounset set -o pipefail set -x REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} if [[ -z "${OUTPUT_DIR:-}" ]]; then OUTPUT_DIR="${REPO_ROOT}/.build/k8s-ai-bench" mkdir -p "${OUTPUT_DIR}" fi echo "Writing results to ${OUTPUT_DIR}" BINDIR="${REPO_ROOT}/.build/bin" mkdir -p "${BINDIR}" curl -sSL https://raw.githubusercontent.com/GoogleCloudPlatform/kubectl-ai/main/install.sh | bash K8S_AI_BENCH_SRC="${REPO_ROOT}/.build/k8s-ai-bench-src" rm -rf "${K8S_AI_BENCH_SRC}" git clone https://github.com/gke-labs/k8s-ai-bench "${K8S_AI_BENCH_SRC}" cd "${K8S_AI_BENCH_SRC}" GOWORK=off go build -o "${BINDIR}/k8s-ai-bench" . "${BINDIR}/k8s-ai-bench" run --agent-bin kubectl-ai --kubeconfig "${KUBECONFIG:-~/.kube/config}" --output-dir "${OUTPUT_DIR}" ${TEST_ARGS:-} ================================================ FILE: dev/ci/presubmits/go-build.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} for f in $(find ${REPO_ROOT} -name go.mod); do cd $(dirname ${f}) go build ./... done ================================================ FILE: dev/ci/presubmits/go-vet.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} for f in $(find ${REPO_ROOT} -name go.mod); do cd $(dirname ${f}) go vet ./... done ================================================ FILE: dev/ci/presubmits/verify-autogen.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} dev/tasks/generate-github-actions.sh changes=$(git status --porcelain) if [[ -n "${changes}" ]]; then echo "FAIL: Changes detected from dev/tasks/generate-github-actions.sh:" git diff | head -n60 echo "${changes}" exit 1 fi ================================================ FILE: dev/ci/presubmits/verify-format.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} dev/tasks/format.sh changes=$(git status --porcelain) if [[ -n "${changes}" ]]; then echo "FAIL: Changes detected from dev/tasks/format.sh:" git diff | head -n60 echo "${changes}" exit 1 fi ================================================ FILE: dev/ci/presubmits/verify-gomod.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} dev/tasks/gomod.sh changes=$(git status --porcelain) if [[ -n "${changes}" ]]; then echo "FAIL: Changes detected from dev/tasks/gomod.sh:" git diff | head -n60 echo "${changes}" exit 1 fi ================================================ FILE: dev/ci/presubmits/verify-mocks.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} if ! command -v mockgen &> /dev/null; then echo "mockgen not found, installing..." go install go.uber.org/mock/mockgen@latest fi # We run generate to see if it creates any diffs. go generate ./internal/mocks if ! git diff --quiet --exit-code -- internal/mocks; then echo "Mocks are stale. Commit the changes to the generated files." exit 1 fi echo "Mocks are up to date." ================================================ FILE: dev/tasks/build-images ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT}/ cd "${SRC_DIR}" if [[ -z "${IMAGE_PREFIX:-}" ]]; then IMAGE_PREFIX="" fi echo "Building images with prefix ${IMAGE_PREFIX}" if [[ -z "${TAG:-}" ]]; then TAG=latest fi if [[ -z "${BUILDX_ARGS:-}" ]]; then BUILDX_ARGS="--load" fi if [[ -z "${DOCKER:-}" ]]; then DOCKER="docker" fi echo "Using container tool: ${DOCKER}" if [[ "${DOCKER}" == "podman" ]]; then # Podman doesn't support buildx, use regular build command if [[ "${BUILDX_ARGS}" == "--push" ]]; then # For podman, build and then push separately echo "Building with podman..." ${DOCKER} build \ --platform linux/amd64 \ -f images/kubectl-ai/Dockerfile \ -t ${IMAGE_PREFIX}kubectl-ai:${TAG} \ . echo "Pushing with podman..." ${DOCKER} push ${IMAGE_PREFIX}kubectl-ai:${TAG} else # For --load or other args, just build ${DOCKER} build \ --platform linux/amd64 \ -f images/kubectl-ai/Dockerfile \ -t ${IMAGE_PREFIX}kubectl-ai:${TAG} \ . fi else # Use docker buildx for docker # Specify platform to avoid multi-arch build requirements ${DOCKER} buildx build ${BUILDX_ARGS} \ --platform linux/amd64 \ -f images/kubectl-ai/Dockerfile \ -t ${IMAGE_PREFIX}kubectl-ai:${TAG} \ --progress=plain . fi ================================================ FILE: dev/tasks/demo.md ================================================ # steps to produce the demo gif 1. Use [asciinema](asciinema.org) to record a screencast. 2. Use agg to produce the gif ```shell agg .github/kubectl-ai.cast --speed 1.3 --idle-time-limit 1 --theme monokai --font-size 24 --cols 100 --rows 25 .github/kubectl-ai.gif ``` ================================================ FILE: dev/tasks/deploy-to-gke ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT} cd "${SRC_DIR}" # Get GCP project ID if [[ -z "${GCP_PROJECT_ID:-}" ]]; then GCP_PROJECT_ID=$(gcloud config get project) fi echo "Using GCP_PROJECT_ID=${GCP_PROJECT_ID}" # Build kubectl command args based on available configuration KUBECTL_ARGS="" if [[ -n "${KUBECONFIG:-}" ]]; then echo "Using KUBECONFIG: ${KUBECONFIG}" else echo "Using default kubeconfig (~/.kube/config)" fi if [[ -n "${KUBE_CONTEXT:-}" ]]; then KUBECTL_ARGS="--context=${KUBE_CONTEXT}" echo "Using kube context: ${KUBE_CONTEXT}" elif [[ -z "${KUBECONFIG:-}" ]]; then # Only require KUBE_CONTEXT if KUBECONFIG is not specified echo "Listing GKE clusters in project ${GCP_PROJECT_ID}:" gcloud container clusters list --project=${GCP_PROJECT_ID} echo "" echo "Please set KUBE_CONTEXT to kubectl context to use, or set KUBECONFIG to use a specific kubeconfig file" exit 1 else echo "Using current context from KUBECONFIG" fi if [[ -z "${NAMESPACE:-}" ]]; then NAMESPACE=kubectl-ai echo "Defaulting to namespace: ${NAMESPACE}" fi # Pick a probably-unique tag export TAG=`date +%Y%m%d%H%M%S` # Set up image registry - default to GCR, but allow override if [[ -z "${IMAGE_REGISTRY:-}" ]]; then IMAGE_REGISTRY=gcr.io/${GCP_PROJECT_ID} fi echo "Using image registry: ${IMAGE_REGISTRY}" # Configure authentication for the container registry before building echo "Configuring authentication for ${IMAGE_REGISTRY}" if [[ "${IMAGE_REGISTRY}" == gcr.io/* ]]; then # Configure GCR authentication gcloud auth configure-docker --quiet elif [[ "${IMAGE_REGISTRY}" == *-docker.pkg.dev/* ]]; then # Configure Artifact Registry authentication gcloud auth configure-docker ${IMAGE_REGISTRY%%/*} --quiet fi # Configure podman authentication before building if needed if [[ "${DOCKER:-docker}" == "podman" ]]; then echo "Using podman: configuring registry authentication" # For podman, we need to configure the auth helper if [[ "${IMAGE_REGISTRY}" == gcr.io/* ]]; then echo "$(gcloud auth print-access-token)" | podman login -u oauth2accesstoken --password-stdin gcr.io elif [[ "${IMAGE_REGISTRY}" == *-docker.pkg.dev/* ]]; then echo "$(gcloud auth print-access-token)" | podman login -u oauth2accesstoken --password-stdin ${IMAGE_REGISTRY%%/*} fi fi # Build the image echo "Building images" export IMAGE_PREFIX=${IMAGE_REGISTRY}/ if [[ -n "${DOCKER:-}" ]]; then echo "Using container tool: ${DOCKER}" export DOCKER fi # For GKE, we need to push images, so use --push instead of --load BUILDX_ARGS=--push dev/tasks/build-images KUBECTL_AI_IMAGE="${IMAGE_PREFIX}kubectl-ai:${TAG}" echo "Built and pushed image: ${KUBECTL_AI_IMAGE}" # Create the namespace if it doesn't exist echo "Creating namespace: ${NAMESPACE}" kubectl create namespace ${NAMESPACE} ${KUBECTL_ARGS} --dry-run=client -oyaml | kubectl apply ${KUBECTL_ARGS} --server-side -f - # Note: No secret needed for Vertex AI - uses Workload Identity for authentication # Create a cluster role binding so kubectl can "see" the current cluster # For production GKE, consider using more restrictive permissions echo "Creating cluster role binding as view" cat < clustername) KIND_CLUSTER_NAME="${KUBE_CONTEXT#kind-}" else # Try to find any available kind cluster AVAILABLE_CLUSTERS=$(kind get clusters 2>/dev/null || echo "") if [[ -n "${AVAILABLE_CLUSTERS}" ]]; then KIND_CLUSTER_NAME=$(echo "${AVAILABLE_CLUSTERS}" | head -n1) echo "Auto-detected kind cluster: ${KIND_CLUSTER_NAME}" else echo "ERROR: No kind clusters found. Please create one with 'kind create cluster'" exit 1 fi fi # Load the image into kind echo "Loading images into kind cluster '${KIND_CLUSTER_NAME}': ${KUBECTL_AI_IMAGE}" if [[ "${DOCKER:-docker}" == "podman" ]]; then # For podman, we need to save and load the image via archive echo "Using podman: saving image to archive first" podman save ${KUBECTL_AI_IMAGE} -o /tmp/kubectl-ai-image.tar kind load image-archive /tmp/kubectl-ai-image.tar --name ${KIND_CLUSTER_NAME} rm -f /tmp/kubectl-ai-image.tar else # For docker, use the standard approach kind load docker-image ${KUBECTL_AI_IMAGE} --name ${KIND_CLUSTER_NAME} fi # Create the namespace if it doesn't exist echo "Creating namespace: ${NAMESPACE}" kubectl create namespace ${NAMESPACE} ${KUBECTL_ARGS} --dry-run=client -oyaml | kubectl apply ${KUBECTL_ARGS} --server-side -f - # Create the secret if it doesn't exist, # including the GEMINI_API_KEY environment variable if set. # (This is for kind, on a GKE cluster, we probably want to use Workload Identity instead) echo "Creating secret: kubectl-ai" cat < ${REPO_ROOT}/.github/workflows/ci-presubmit.yaml <> ${REPO_ROOT}/.github/workflows/ci-presubmit.yaml <> ${REPO_ROOT}/.github/workflows/ci-presubmit.yaml <<'EOF' concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} cancel-in-progress: true EOF ================================================ FILE: dev/tasks/gomod.sh ================================================ #!/usr/bin/env bash # Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" cd ${REPO_ROOT} # retry # Retries a command on failure with a fixed delay between attempts. # Useful for transient Go proxy / network errors (e.g. HTTP/2 INTERNAL_ERROR). retry() { local max=$1; shift local delay=$1; shift local attempt=1 until "$@"; do if (( attempt >= max )); then echo "ERROR: command failed after ${max} attempts: $*" >&2 return 1 fi echo "WARNING: attempt ${attempt}/${max} failed, retrying in ${delay}s: $*" >&2 sleep "${delay}" (( attempt++ )) done } for f in $(find ${REPO_ROOT} -name go.mod); do cd $(dirname ${f}) rm go.sum retry 3 5 go mod tidy done ================================================ FILE: docs/bedrock.md ================================================ # AWS Bedrock Provider kubectl-ai supports AWS Bedrock models including Claude Sonnet 4 and Claude 3.7. ## Setup ### AWS Credentials Configure AWS credentials using standard AWS SDK methods: ```bash # Option 1: Environment variables export AWS_ACCESS_KEY_ID="your-access-key" export AWS_SECRET_ACCESS_KEY="your-secret-key" export AWS_REGION="us-east-1" # Option 2: AWS Profile (recommended) export AWS_PROFILE="your-profile-name" export AWS_REGION="us-east-1" # Option 3: Use IAM roles (on EC2/ECS/Lambda) export AWS_REGION="us-east-1" ``` ### Model Configuration ```bash # Optional: Set default model export BEDROCK_MODEL="us.anthropic.claude-3-7-sonnet-20250219-v1:0" ``` ## Supported Models See [AWS Bedrock documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html) for current model availability and regional support. Currently supported: - Claude Sonnet 4: `us.anthropic.claude-sonnet-4-20250514-v1:0` (default) - Claude 3.7 Sonnet: `us.anthropic.claude-3-7-sonnet-20250219-v1:0` ## Usage ```bash # Use default model (Claude Sonnet 4) kubectl-ai --provider bedrock "explain this deployment" # Specify model explicitly kubectl-ai --provider bedrock --model us.anthropic.claude-3-7-sonnet-20250219-v1:0 "help me debug this pod" ``` ## Authentication kubectl-ai uses the standard AWS SDK credential provider chain: 1. Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) 2. AWS credentials file (~/.aws/credentials) 3. AWS config file (~/.aws/config) 4. IAM roles for EC2 instances 5. IAM roles for ECS tasks 6. IAM roles for Lambda functions For more details, see [AWS SDK Go Configuration](https://aws.github.io/aws-sdk-go-v2/docs/configuring-sdk/). ## Region Configuration Bedrock is available in specific AWS regions. Set your region using: ```bash export AWS_REGION="us-east-1" # Primary Bedrock region ``` Alternatively, configure region in `~/.aws/config`: ```ini [default] region = us-east-1 ``` ================================================ FILE: docs/gke-deployment.md ================================================ # Deploying k8-kate to Google Kubernetes Engine ## Prerequisites - A GKE cluster (Standard or Autopilot). - `gcloud` CLI authenticated with `gcloud auth login`. - Local Docker environment (or Cloud Build) capable of building and pushing container images. - [`kubectl` configured](https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl) to talk to your target cluster. - A Gemini API key with access to the model you plan to demo. ## 1. Set your Google Cloud context ```bash export PROJECT_ID="my-gcp-project" export REGION="us-central1" export CLUSTER_NAME="kubectl-ai-demo" gcloud config set project "${PROJECT_ID}" gcloud container clusters get-credentials "${CLUSTER_NAME}" --region "${REGION}" ``` These commands configure both `gcloud` and `kubectl` to operate on the cluster that will host `kubectl-ai`. ## 2. Build and push the kubectl-ai image Pick an Artifact Registry repository or another container registry that your cluster can pull from. The snippet below creates an Artifact Registry repo (if needed), builds the image locally, and pushes it. ```bash # Create an Artifact Registry (skip if you already have one) gcloud artifacts repositories create kubectl-ai \ --location="${REGION}" \ --repository-format=DOCKER \ --description="kubectl-ai demo images" # Configure Docker to authenticate to Artifact Registry gcloud auth configure-docker "${REGION}"-docker.pkg.dev # Build and push the container IMAGE="${REGION}-docker.pkg.dev/${PROJECT_ID}/kubectl-ai/kubectl-ai:latest" docker build -t "${IMAGE}" -f images/kubectl-ai/Dockerfile . docker push "${IMAGE}" ``` ## 3. Prepare cluster namespaces and RBAC Create the namespaces and RBAC that the hosted agent requires: ```bash # Sandbox namespace + RBAC (creates `computer` namespace, service account, and reader roles) kubectl apply -f k8s/sandbox/all-in-one.yaml ``` The sandbox manifest provisions the `computer` namespace and the `normal-user` service account used for sandbox pods. The `kubectl-ai-gke.yaml` manifest will create the `kubectl-ai` namespace automatically. ## 4. Configure the deployment manifest Copy `k8s/kubectl-ai-gke.yaml` to a working file and edit the following sections: 1. **Container image** – replace the `REPLACE_WITH_YOUR_IMAGE` 2. **Gemini API key** – change `REPLACE_WITH_YOUR_GEMINI_API_KEY` to the key you obtained from Google AI Studio Review the RBAC objects in the manifest and adjust them if your security posture requires tighter permissions. ## 5. Deploy kubectl-ai Apply the updated manifest to your cluster: ```bash kubectl apply -f kubectl-ai-gke.yaml ``` Kubernetes creates the Deployment, ServiceAccount, RBAC bindings, and Service for the hosted agent. You can watch the rollout with: ```bash kubectl get pods -n kubectl-ai kubectl describe pod -n kubectl-ai -l app=kubectl-ai | grep -i image ``` ## 6. Access the hosted web UI Port-forward the Service locally to interact with the hosted UI: ```bash kubectl port-forward svc/kubectl-ai -n kubectl-ai 8080:80 ``` Then open [http://localhost:8080](http://localhost:8080) in your browser. Each browser session can create, rename, and delete conversations, and messages stream in real time via Server-Sent Events. If you prefer to expose the UI via an external Load Balancer, replace the Service type in the manifest with `LoadBalancer` and configure the appropriate firewall rules. ## 7. Verify sandboxed tool execution When the UI creates a conversation, the agent launches a sandbox pod in the `computer` namespace. You can confirm sandbox activity with: ```bash kubectl get pods -n computer ``` Pods named `kubectl-ai-sandbox-*` indicate that commands are running inside isolated helper containers. Asking the agent to execute commands such as `uname -a` should produce output that matches the sandbox image (for example `bitnami/kubectl`). ## 8. Cleanup Remove the deployment and sandbox resources when you are done: ```bash kubectl delete -f kubectl-ai-gke.yaml kubectl delete namespace kubectl-ai kubectl delete -f k8s/sandbox/all-in-one.yaml ``` If you no longer need the Artifact Registry repository or pushed image, delete them using `gcloud artifacts repositories delete` and `gcloud artifacts docker images delete`. ================================================ FILE: docs/mcp-client.md ================================================ # kubectl-ai MCP Client Integration ## Multi-Server Orchestration for Security Automation The MCP (Model Context Protocol) Client feature enables kubectl-ai to coordinate multiple specialized tools through natural language commands. This integration demonstrates automated security workflows that combine RBAC scanning with email reporting. **Problem**: Traditional security audits require manual execution of multiple tools, data correlation, and report distribution—a time-consuming process prone to human error. **Solution**: Single command orchestration across multiple MCP servers: ```bash kubectl-ai --mcp-client --quiet "scan rbac and send urgent report to incident-team@company.com from sender@company.com" ``` **Architecture Components:** - **kubectl-ai**: Central orchestrator interpreting natural language commands - **Permiflow**: RBAC security scanning and analysis - **Resend**: Automated email delivery service - **Additional servers**: Documentation, reasoning, and extensible integrations ## Workflow Sequence Diagram ```mermaid sequenceDiagram participant User participant kubectl-ai as kubectl-ai
(MCP Client) participant Permiflow as Permiflow
(MCP Server) participant K8s as Kubernetes
Cluster participant Resend as Resend
(MCP Server) participant Email as Email
Recipient User->>kubectl-ai: "scan rbac and send report to admin@company.com" kubectl-ai->>Permiflow: scan_rbac() Permiflow->>K8s: Query RBAC policies K8s-->>Permiflow: Return roles, bindings, permissions Permiflow->>Permiflow: Analyze security risks Permiflow-->>kubectl-ai: Security findings report kubectl-ai->>kubectl-ai: Format report for email kubectl-ai->>Resend: send_email(to, from, subject, content) Resend->>Email: Deliver formatted security report Email-->>User: Email confirmation kubectl-ai-->>User: "✅ RBAC scan completed and report sent" ``` ## Execution Flow The command execution follows this sequence: 1. **kubectl-ai** parses the natural language request 2. [**Permiflow**](https://github.com/tutran-se/permiflow) performs comprehensive RBAC analysis across cluster resources 3. [**Resend**](https://github.com/resend/mcp-send-email) formats and delivers the security report via email **Extensibility**: The architecture supports additional MCP servers for Slack notifications, Jira ticket creation, compliance databases, and custom integrations. ## Configuration and Setup ### MCP Server Configuration Configure the MCP servers in `~/.config/kubectl-ai/mcp.yaml`: ```yaml servers: - name: resend command: node args: - "~/mcp-send-email/build/index.js" env: RESEND_API_KEY: "api-key-here" - name: permiflow url: http://localhost:8080/mcp ``` ### Quick Start ```bash # 1. Start the Permiflow MCP server permiflow mcp --transport http --http-port 8080 # 2. Execute kubectl-ai with MCP client enabled kubectl-ai --mcp-client --quiet "scan rbac and send report to admin@company.com from sec@company.com" ``` ## Automation Use Cases ### Scheduled Security Monitoring Implement automated daily security scans using cron: ```bash # Daily RBAC audit at 9 AM 0 9 * * * kubectl-ai --mcp-client --quiet "scan rbac and send daily report to admin@company.com from sec@company.com" ``` ### Incident Response Execute immediate security assessments during incidents: ```bash kubectl-ai --mcp-client --quiet "scan rbac for production namespace and send urgent report to incident-team@company.com from sec@company.com" ``` ## Usage Examples ### Interactive Mode Launch kubectl-ai in interactive mode for exploratory analysis: ```bash kubectl-ai --mcp-client >>> "scan rbac and send report to admin@company.com" >>> "analyze RBAC for kubeflow namespace" >>> "show me the most dangerous permissions in production" >>> "which service accounts can access secrets across namespaces?" ``` ### Direct Commands Execute specific security queries directly: ```bash kubectl-ai --mcp-client "show wildcard permissions and suggest fixes" ``` ## Extended Integration ### Additional MCP Servers Expand the automation capabilities by adding specialized servers: ```yaml servers: - name: slack-notifier url: "https://slack-mcp.company.com/mcp" - name: jira-tickets url: "https://jira-mcp.company.com/mcp" - name: trivy-scanner command: npx args: ["-y", "@aquasecurity/trivy-mcp"] ``` ### Advanced Workflows **Multi-Channel Incident Response:** ```bash "scan rbac, create jira ticket, email security team, post to slack" ``` **Compliance Automation:** ```bash "scan vulnerabilities, update compliance database, email leadership" ``` ## Benefits - **Unified Interface**: Single natural language interface for multiple tools - **Automation**: Reduces manual security audit processes - **Consistency**: Standardized security scanning and reporting - **Extensibility**: Modular architecture supports additional integrations - **Efficiency**: Rapid security assessment and stakeholder notification ================================================ FILE: docs/mcp-server.md ================================================ # kubectl-ai MCP Server kubectl-ai can run as an MCP (Model Context Protocol) server, exposing kubectl-ai tools to other MCP clients. The server can run in two modes: 1. **Built-in tools only**: Exposes only kubectl-ai's native tools 2. **External tool discovery**: Additionally discovers and exposes tools from other MCP servers ## Quick Start ### Basic MCP Server (Built-in tools only) Start the MCP server with only kubectl-ai's built-in tools: ```bash kubectl-ai --mcp-server ``` ### Enhanced MCP Server (With external tool discovery) Start the MCP server with external MCP tool discovery enabled: ```bash kubectl-ai --mcp-server --external-tools ``` ### Expose an HTTP Endpoint for MCP Clients Run the server with the streamable HTTP transport to serve compatible MCP clients (including kubectl-ai MCP client mode) over HTTP: ```bash kubectl-ai --mcp-server --mcp-server-mode streamable-http --http-port 9080 ``` This listens on `http://localhost:9080/mcp` by default. ## Configuration When `--external-tools` is enabled, the enhanced MCP server will automatically discover and expose tools from configured MCP servers. You can configure MCP servers using the standard MCP client configuration file. ### Example MCP Configuration Create `~/.config/kubectl-ai/mcp.yaml`: ```yaml servers: filesystem: command: "npx" args: [ "-y", "@modelcontextprotocol/server-filesystem", "/path/to/allowed/files", ] brave-search: command: "npx" args: ["-y", "@modelcontextprotocol/server-brave-search"] env: BRAVE_API_KEY: "your-api-key" ``` ## Features ### Tool Aggregation When external tool discovery is enabled with `--external-tools`, the kubectl-ai MCP server acts as a **tool aggregator**, providing: - All kubectl-ai built-in tools (kubectl, cluster analysis, etc.) - Tools from external MCP servers (filesystem, web search, etc.) - Unified interface for all tools through a single MCP endpoint ### Graceful Degradation The server handles external MCP connection failures gracefully: - If external MCP servers are unavailable, the server continues with built-in tools only - Individual tool failures don't affect the overall server operation - Clear logging for troubleshooting connection issues ### Example Usage in Claude Desktop Configure Claude Desktop to use kubectl-ai as an MCP server: **Basic usage (built-in tools only):** ```json { "mcpServers": { "kubectl-ai": { "command": "kubectl-ai", "args": ["--mcp-server"] } } } ``` **Enhanced usage (with external tools):** ```json { "mcpServers": { "kubectl-ai": { "command": "kubectl-ai", "args": ["--mcp-server", "--external-tools"] } } } ``` ## Available Tools ### Built-in Tools kubectl-ai provides the following native tools: - `bash`: Executes a bash command. Use this tool only when you need to execute a shell command. - `kubectl`: Executes a kubectl command against the user's Kubernetes cluster. Use this tool only when you need to query or modify the state of the user's Kubernetes cluster. ### External Tools (when `--external-tools` is enabled) Additional tools are available depending on the configured MCP servers: - **Filesystem tools**: Read/write files, list directories - **Web search tools**: Search the internet for information - **Database tools**: Query databases - **API tools**: Interact with external APIs - **Custom tools**: Any MCP-compatible tools ## Command Line Options | Flag | Default | Description | | ------------------- | ---------------- | ---------------------------------------------------------------------- | | `--mcp-server` | `false` | Run in MCP server mode | | `--external-tools` | `false` | Discover and expose external MCP tools (requires --mcp-server) | | `--kubeconfig` | `~/.kube/config` | Path to kubeconfig file | | `--mcp-server-mode` | `stdio` | Transport for the MCP server (`stdio` or `streamable-http`) | | `--http-port` | `9080` | Port for the HTTP endpoint when using `streamable-http` modes | ## Architecture ```txt ┌─────────────────┐ ┌───────────────────┐ ┌─────────────────┐ │ MCP Client │───▶│ kubectl-ai Server │───▶│ External Tools │ │ (Claude, etc.) │ │ │ │ (filesystem, │ │ │ │ ┌───────────────┐ │ │ web search, │ │ │ │ │ Built-in │ │ │ etc.) │ │ │ │ │ kubectl tools │ │ │ │ │ │ │ └───────────────┘ │ │ │ └─────────────────┘ └───────────────────┘ └─────────────────┘ ``` The kubectl-ai MCP server acts as both: - An **MCP Server** (exposing tools to clients) - An **MCP Client** (consuming tools from other servers, when `--external-tools` is enabled) This creates a powerful tool aggregation pattern where kubectl-ai becomes a central hub for both Kubernetes operations and general-purpose tools. ## Troubleshooting ### External Tools Not Available If external tools aren't appearing: 1. Ensure you're using both `--mcp-server` and `--external-tools` flags 2. Check MCP configuration file exists and is valid 3. Verify external MCP servers are working independently 4. Check kubectl-ai logs for connection errors 5. Try running with external tools disabled to isolate issues ### Performance Considerations - Tool discovery adds startup time (usually 2-3 seconds) when `--external-tools` is enabled - Each external tool call has network overhead - Consider running without `--external-tools` for faster startup if external tools aren't needed ### Debugging Enable verbose logging to troubleshoot: ```bash kubectl-ai --mcp-server --external-tools -v=2 ``` This will show: - MCP server connection attempts - Tool discovery results - Tool call routing decisions ================================================ FILE: docs/mocking.md ================================================ # Mocking in kubectl-ai ## Gomock developer workflow We use [gomock](https://github.com/uber-go/mock) to mock external dependencies. All mocks and generated files live under `internal/mocks/`. - **Everyday commands** - Regenerate mocks after changing interfaces or adding new ones: `make generate` - Verify nothing is stale (locally or in CI): `make verify-mocks` - Run tests: `go test ./...` - **Generator install** (if you don’t have it yet): `go install go.uber.org/mock/mockgen@latest` - **What `make generate` does** Runs `go generate ./internal/mocks`. Note: `go generate` is **not** part of `go build/test`; commit generated mocks. - **Add a new mock** 1. Add a `go:generate` line in `internal/mocks/generate.go`, e.g.: ```go //go:generate mockgen -destination=tools_mock.go -package=mocks // github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools Tool ``` 2. Run `make generate` and import the mocks in tests: ```go ctrl := gomock.NewController(t) defer ctrl.Finish() llm := mocks.NewMockClient(ctrl) // example llm.EXPECT().NewChat(gomock.Any()).AnyTimes() ``` ## When and when not to use gomock **Use gomock for:** - **External boundaries / side effects**: `gollm.Client`, `gollm.Chat`, `pkg/tools.Tool`, network/IO, anything slow or flaky. - **Behavioral checks**: asserting specific calls/arguments or injecting failures/timeouts. **Prefer fakes/in‑memory over mocks for:** - **Stateful components with an in‑memory impl** (e.g., session/message store). Don’t mock storage if an in‑memory version exists. - **Pure functions / simple value types**—call them directly. **Good practices:** - Keep expectations minimal—assert only what matters. Use `gomock.Any()` and `AnyTimes()`/`MinTimes(1)` where exact call counts don’t matter. - Centralize `mockgen` directives in `internal/mocks/generate.go`. - **If an interface changes**: run `make generate`, fix compile errors in tests (signatures/matchers), update/remove `go:generate` lines if package paths or names changed, and commit the regenerated mocks. ================================================ FILE: docs/tool-samples/argocd.yaml ================================================ - name: argocd description: "A declarative, GitOps continuous delivery tool for Kubernetes. Use it to manage application deployments from Git repositories." command: "argocd" command_desc: | The argocd command-line interface. Core subcommands and usage patterns: - `argocd login [flags]`: Log in to an Argo CD server. This is required before most other commands. - `argocd app list [flags]`: List all applications managed by Argo CD. - `argocd app get [flags]`: Get detailed information about a specific application. - `argocd app sync [flags]`: Sync an application to its target state defined in the Git repository. - `argocd app history [flags]`: View the deployment history of an application. - `argocd app rollback `: Roll back an application to a previous deployed state. - `argocd app create [flags]`: Create a new application. - `argocd app delete `: Delete an application. Use `argocd --help` or `argocd --help` for full syntax and available flags. ================================================ FILE: docs/tool-samples/gcloud.yaml ================================================ - name: gcloud description: "The gcloud command-line tool is the primary CLI for Google Cloud. Use it to manage Google Cloud resources, including Google Kubernetes Engine (GKE) clusters, virtual machines, and networking." command: "gcloud" command_desc: | The gcloud CLI manages authentication, local configuration, developer workflow, and interactions with Google Cloud APIs. For Kubernetes-related tasks, the `gcloud container clusters` command group is the most relevant. Core subcommands and usage patterns: - `gcloud container clusters list [flags]`: List all GKE clusters in the configured project and zone/region. - `gcloud container clusters describe [flags]`: Get detailed information about a specific GKE cluster. - `gcloud container clusters get-credentials [flags]`: Fetch cluster endpoint and auth data and configure kubectl to use the cluster. This is essential for connecting to a GKE cluster. - `gcloud container clusters create [flags]`: Create a new GKE cluster. - `gcloud container clusters delete [flags]`: Delete an existing GKE cluster. You can set the default project, region, and zone using: - `gcloud config set project ` - `gcloud config set compute/region ` - `gcloud config set compute/zone ` Use `gcloud --help` or `gcloud --help` for full syntax and available flags. ================================================ FILE: docs/tool-samples/gh.yaml ================================================ - name: gh description: "The official GitHub command-line tool. Use it to interact with GitHub repositories, pull requests, issues, actions, and more, directly from the terminal." command: "gh" command_desc: | The gh command-line interface for GitHub. Core subcommands and usage patterns: - `gh auth login`: Authenticate with a GitHub host. This is required before most other commands. - `gh repo view [repo]`: View a repository. - `gh pr list`: List pull requests in the current repository. - `gh pr view `: View a specific pull request. - `gh pr checkout `: Check out a pull request locally. - `gh pr create [flags]`: Create a new pull request. - `gh pr merge [flags]`: Merge a pull request. - `gh issue list`: List issues in the current repository. - `gh issue view `: View a specific issue. - `gh workflow list`: List workflows in the current repository. - `gh workflow run [flags]`: Trigger a workflow run. Use `gh --help` or `gh --help` for full syntax and available flags. ================================================ FILE: docs/tool-samples/kustomize.yaml ================================================ - name: kustomize description: "A tool to customize Kubernetes resource configurations. Use it to render and apply declarative configurations from a directory containing a kustomization.yaml file." command: "kustomize" command_desc: | The kustomize command-line interface. Core subcommands and usage patterns: - `kustomize build `: Prints the customized resources to standard output. This is useful for inspecting the final configuration before applying it. - `kustomize build | kubectl apply -f -`: A common pattern to apply the output directly to the cluster. Note: `kubectl apply -k ` is a shorthand for the pipe command above and is often preferred. ================================================ FILE: docs/tools.md ================================================ # Custom Tools for kubectl-ai `kubectl-ai` leverages LLMs to suggest and execute Kubernetes operations using a set of powerful tools. It comes with built-in tools like `kubectl` and `bash`. The `kubectl-ai` assistant can be extended with custom tools to interact with various command-line interfaces (CLIs) beyond `kubectl`. This allows the AI to perform a wider range of tasks related to infrastructure management, CI/CD, and more. This document outlines how you can add custom tools by detailing the steps and providing samples. This document also outlines the available tools, their locations, and how to use them. ## Adding Custom Tools Custom tools can be added by following these two steps: - describing or templating the tool through YAML file - enabling the tool in the configuration file by pointing the **--custom-tools-config** to this file / directory ## Describing the Tool in YAML file A custom tool can be described by providing the following four pieces of information: - **name**: name of the tool - **description**: "A clear description that helps the LLM understand when to use this tool." - **command** : "your_command" # For example: 'gcloud' or 'gcloud container clusters' - **command_desc**: "Detailed information for the LLM, including command syntax and usage examples." Samples are provided in the `pkg/tools/samples` directory. Below is a sample for the `kustomize` tool: ```yaml - name: kustomize description: "A tool to customize Kubernetes resource configurations. Use it to render and apply declarative configurations from a directory containing a kustomization.yaml file." command: "kustomize" command_desc: | The kustomize command-line interface. Core subcommands and usage patterns: - `kustomize build `: Prints the customized resources to standard output. This is useful for inspecting the final configuration before applying it. - `kustomize build | kubectl apply -f -`: A common pattern to apply the output directly to the cluster. Note: `kubectl apply -k ` is a shorthand for the pipe command above and is often preferred. ``` ## Enabling the Custom Tool To enable the custom tools, you must point `kubectl-ai` to the directory containing the tool configuration YAML files using the `--custom-tools-config` flag. `kubectl-ai` can pick up a single YAML file (e.g., `tools.yaml`) containing all the tool descriptions or multiple individual YAML files when pointed to a directory containing them. This example uses multiple YAML files located in a single directory. In case, you don't want to use a tool (that is provided in samples), just move the file out of the directory into some other location & restart `kubectl-ai`. ### Running from a Local Binary When running the `kubectl-ai` binary directly, provide the path to your local tools directory. ```sh ./kubectl-ai --custom-tools-config= "your prompt here" ``` ### Running with Docker Image When using the Docker image, you can either use the tools baked into the image or mount your own custom directory. #### Using Built-in Tools The official Docker image includes the default tool configurations. You can enable them by pointing to the internal path. ```sh docker run --rm -it your-kubectl-ai-image:latest \ --custom-tools-config=/etc/kubectl-ai/tools \ "list all pull requests on GitHub" ``` #### Using a Local Tools Directory To use a custom set of tools from your local machine, mount the directory into the container and point the flag to the mounted path. This is useful for developing and testing new tools. ```sh docker run --rm -it \ -v /path/to/your/local/tools:/my-custom-tools \ your-kubectl-ai-image:latest \ --custom-tools-config=/my-custom-tools \ "your prompt here" ``` ## Sample Custom Tools The following sample custom tools are configured by default. | Tool | Description | YAML File | | :--------------------------------------------------------- | :-------------------------------------------------------------- | :------------------------------------------------------ | | Argo CD (`argocd`) | A declarative, GitOps continuous delivery tool for Kubernetes. | [argocd.yaml](./tool-samples/argocd.yaml) | | GitHub CLI (`gh`) | The official command-line tool to interact with GitHub. | [gh.yaml](./tool-samples/gh.yaml) | | Google Cloud CLI (`gcloud`) | The primary CLI for managing Google Cloud resources. | [gcloud.yaml](./tool-samples/gcloud.yaml) | | Kustomize (`kustomize`) | A tool to customize Kubernetes resource configurations. | [kustomize.yaml](./tool-samples/kustomize.yaml) | ================================================ FILE: go.mod ================================================ module github.com/GoogleCloudPlatform/kubectl-ai go 1.24.0 toolchain go1.24.3 // Needed for multiple go modules in one repo replace github.com/GoogleCloudPlatform/kubectl-ai/gollm => ./gollm require ( github.com/GoogleCloudPlatform/kubectl-ai/gollm v0.0.0-00010101000000-000000000000 github.com/charmbracelet/bubbles v0.21.0 github.com/charmbracelet/bubbletea v1.3.5 github.com/charmbracelet/glamour v0.10.0 github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 github.com/chzyer/readline v1.5.1 github.com/google/uuid v1.6.0 github.com/mark3labs/mcp-go v0.41.1 github.com/spf13/cobra v1.9.1 github.com/spf13/pflag v1.0.6 go.uber.org/mock v0.6.0 golang.org/x/sync v0.16.0 golang.org/x/term v0.33.0 k8s.io/api v0.34.2 k8s.io/apimachinery v0.34.2 k8s.io/client-go v0.34.2 k8s.io/klog/v2 v2.130.1 mvdan.cc/sh/v3 v3.11.0 sigs.k8s.io/yaml v1.6.0 ) require ( cloud.google.com/go v0.118.3 // indirect cloud.google.com/go/auth v0.15.0 // indirect cloud.google.com/go/compute/metadata v0.6.0 // indirect github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.7.2 // indirect github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 // indirect github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0 // indirect github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 // indirect github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.7.0 // indirect github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.2.0 // indirect github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2 // indirect github.com/alecthomas/chroma/v2 v2.14.0 // indirect github.com/anthropics/anthropic-sdk-go v1.26.0 // indirect github.com/atotto/clipboard v0.1.4 // indirect github.com/aws/aws-sdk-go-v2 v1.36.6 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.11 // indirect github.com/aws/aws-sdk-go-v2/config v1.29.18 // indirect github.com/aws/aws-sdk-go-v2/credentials v1.17.71 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.33 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.37 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.37 // indirect github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 // indirect github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.31.1 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.4 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.18 // indirect github.com/aws/aws-sdk-go-v2/service/sso v1.25.6 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.34.1 // indirect github.com/aws/smithy-go v1.22.4 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/aymerick/douceur v0.2.0 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.1 // indirect github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect github.com/charmbracelet/x/ansi v0.8.0 // indirect github.com/charmbracelet/x/cellbuf v0.0.13 // indirect github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf // indirect github.com/charmbracelet/x/term v0.2.1 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/dlclark/regexp2 v1.11.4 // indirect github.com/emicklei/go-restful/v3 v3.12.2 // indirect github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/fxamacker/cbor/v2 v2.9.0 // indirect github.com/go-logr/logr v1.4.2 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/go-openapi/jsonpointer v0.21.0 // indirect github.com/go-openapi/jsonreference v0.20.2 // indirect github.com/go-openapi/swag v0.23.0 // indirect github.com/gogo/protobuf v1.3.2 // indirect github.com/golang-jwt/jwt/v5 v5.2.2 // indirect github.com/google/gnostic-models v0.7.0 // indirect github.com/google/go-cmp v0.7.0 // indirect github.com/google/s2a-go v0.1.9 // indirect github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect github.com/googleapis/gax-go/v2 v2.14.1 // indirect github.com/gorilla/css v1.0.1 // indirect github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/kylelemons/godebug v1.1.0 // indirect github.com/lucasb-eyer/go-colorful v1.2.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-localereader v0.0.1 // indirect github.com/mattn/go-runewidth v0.0.16 // indirect github.com/microcosm-cc/bluemonday v1.0.27 // indirect github.com/moby/spdystream v0.5.0 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect github.com/muesli/reflow v0.3.0 // indirect github.com/muesli/termenv v0.16.0 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect github.com/ollama/ollama v0.6.5 // indirect github.com/openai/openai-go v1.12.0 // indirect github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect github.com/pkg/errors v0.9.1 // indirect github.com/rivo/uniseg v0.4.7 // indirect github.com/sahilm/fuzzy v0.1.1 // indirect github.com/spf13/cast v1.7.1 // indirect github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/x448/float16 v0.8.4 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect github.com/yuin/goldmark v1.7.8 // indirect github.com/yuin/goldmark-emoji v1.0.5 // indirect go.opentelemetry.io/auto/sdk v1.1.0 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 // indirect go.opentelemetry.io/otel v1.34.0 // indirect go.opentelemetry.io/otel/metric v1.34.0 // indirect go.opentelemetry.io/otel/trace v1.34.0 // indirect go.yaml.in/yaml/v2 v2.4.2 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect golang.org/x/crypto v0.40.0 // indirect golang.org/x/net v0.41.0 // indirect golang.org/x/oauth2 v0.30.0 // indirect golang.org/x/sys v0.34.0 // indirect golang.org/x/text v0.27.0 // indirect golang.org/x/time v0.10.0 // indirect google.golang.org/genai v1.8.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2 // indirect google.golang.org/grpc v1.70.0 // indirect google.golang.org/protobuf v1.36.5 // indirect gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect k8s.io/kube-openapi v0.0.0-20250710124328-f3f2b991d03b // indirect k8s.io/utils v0.0.0-20250604170112-4c0f3b243397 // indirect sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8 // indirect sigs.k8s.io/randfill v1.0.0 // indirect sigs.k8s.io/structured-merge-diff/v6 v6.3.0 // indirect ) ================================================ FILE: go.sum ================================================ cloud.google.com/go v0.118.3 h1:jsypSnrE/w4mJysioGdMBg4MiW/hHx/sArFpaBWHdME= cloud.google.com/go v0.118.3/go.mod h1:Lhs3YLnBlwJ4KA6nuObNMZ/fCbOQBPuWKPoE0Wa/9Vc= cloud.google.com/go/auth v0.15.0 h1:Ly0u4aA5vG/fsSsxu98qCQBemXtAtJf+95z9HK+cxps= cloud.google.com/go/auth v0.15.0/go.mod h1:WJDGqZ1o9E9wKIL+IwStfyn/+s59zl4Bi+1KQNVXLZ8= cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I= cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg= github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.7.2 h1:+hDUZnYHHoXu05iXiJcL53MZW7raZZejB8ZtzVW7yyc= github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.7.2/go.mod h1:49PyorVrwk6G+e8Vghvn7EkAS6wSPdXEu5a8iW2/vC8= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 h1:Gt0j3wceWMwPmiazCa8MzMA0MfhmPIz0Qp0FJ6qcM0U= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0/go.mod h1:Ot/6aikWnKWi4l9QB7qVSwa8iMphQNqkWALMoNT3rzM= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0 h1:OVoM452qUFBrX+URdH3VpR299ma4kfom0yB0URYky9g= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0/go.mod h1:kUjrAo8bgEwLeZ/CmHqNl3Z/kPm7y6FKfxxK0izYUg4= github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2 h1:yz1bePFlP5Vws5+8ez6T3HWXPmwOK7Yvq8QxDBD3SKY= github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2/go.mod h1:Pa9ZNPuoNu/GztvBSKk9J1cDJW6vk/n0zLtV4mgd8N8= github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 h1:FPKJS1T+clwv+OLGt13a8UjqeRuh0O4SJ3lUriThc+4= github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1/go.mod h1:j2chePtV91HrC22tGoRX3sGY42uF13WzmmV80/OdVAA= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.7.0 h1:4exaC92+n1FzhSKb5Ghino2XEk3cClUtzvveL1U9YeM= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.7.0/go.mod h1:BkhZrH3JiVTkrTqCeYHOmqReFcZTYEMf8jcFDlrCJLk= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.2.0 h1:UrGzkHueDwAWDdjQxC+QaXHd4tVCkISYE9j7fSSXF8k= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.2.0/go.mod h1:qskvSQeW+cxEE2bcKYyKimB1/KiQ9xpJ99bcHY0BX6c= github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJTmL004Abzc5wDB5VtZG2PJk5ndYDgVacGqfirKxjM= github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1/go.mod h1:tCcJZ0uHAmvjsVYzEFivsRTN00oz5BEsRgQHu5JZ9WE= github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2 h1:oygO0locgZJe7PpYPXT5A29ZkwJaPqcva7BVeemZOZs= github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI= github.com/alecthomas/assert/v2 v2.7.0 h1:QtqSACNS3tF7oasA8CU6A6sXZSBDqnm7RfpLl9bZqbE= github.com/alecthomas/assert/v2 v2.7.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= github.com/alecthomas/chroma/v2 v2.14.0 h1:R3+wzpnUArGcQz7fCETQBzO5n9IMNi13iIs46aU4V9E= github.com/alecthomas/chroma/v2 v2.14.0/go.mod h1:QolEbTfmUHIMVpBqxeDnNBj2uoeI4EbYP4i6n68SG4I= github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc= github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= github.com/anthropics/anthropic-sdk-go v1.26.0 h1:oUTzFaUpAevfuELAP1sjL6CQJ9HHAfT7CoSYSac11PY= github.com/anthropics/anthropic-sdk-go v1.26.0/go.mod h1:qUKmaW+uuPB64iy1l+4kOSvaLqPXnHTTBKH6RVZ7q5Q= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs= github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= github.com/aws/aws-sdk-go-v2 v1.36.6 h1:zJqGjVbRdTPojeCGWn5IR5pbJwSQSBh5RWFTQcEQGdU= github.com/aws/aws-sdk-go-v2 v1.36.6/go.mod h1:EYrzvCCN9CMUTa5+6lf6MM4tq3Zjp8UhSGR/cBsjai0= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.11 h1:12SpdwU8Djs+YGklkinSSlcrPyj3H4VifVsKf78KbwA= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.11/go.mod h1:dd+Lkp6YmMryke+qxW/VnKyhMBDTYP41Q2Bb+6gNZgY= github.com/aws/aws-sdk-go-v2/config v1.29.18 h1:x4T1GRPnqKV8HMJOMtNktbpQMl3bIsfx8KbqmveUO2I= github.com/aws/aws-sdk-go-v2/config v1.29.18/go.mod h1:bvz8oXugIsH8K7HLhBv06vDqnFv3NsGDt2Znpk7zmOU= github.com/aws/aws-sdk-go-v2/credentials v1.17.71 h1:r2w4mQWnrTMJjOyIsZtGp3R3XGY3nqHn8C26C2lQWgA= github.com/aws/aws-sdk-go-v2/credentials v1.17.71/go.mod h1:E7VF3acIup4GB5ckzbKFrCK0vTvEQxOxgdq4U3vcMCY= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.33 h1:D9ixiWSG4lyUBL2DDNK924Px9V/NBVpML90MHqyTADY= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.33/go.mod h1:caS/m4DI+cij2paz3rtProRBI4s/+TCiWoaWZuQ9010= github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.37 h1:osMWfm/sC/L4tvEdQ65Gri5ZZDCUpuYJZbTTDrsn4I0= github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.37/go.mod h1:ZV2/1fbjOPr4G4v38G3Ww5TBT4+hmsK45s/rxu1fGy0= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.37 h1:v+X21AvTb2wZ+ycg1gx+orkB/9U6L7AOp93R7qYxsxM= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.37/go.mod h1:G0uM1kyssELxmJ2VZEfG0q2npObR3BAkF3c1VsfVnfs= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 h1:bIqFDwgGXXN1Kpp99pDOdKMTTb5d2KyU5X/BZxjOkRo= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3/go.mod h1:H5O/EsxDWyU+LP/V8i5sm8cxoZgc2fdNR9bxlOFrQTo= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.31.1 h1:JDLT1baDmioiZKa2bZ6J82/Zwfv9cSAjr+LyF47TPYw= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.31.1/go.mod h1:FvbGcqrU4sC3qjrAKK3FzOmBoucDJF2dXsKVvAbGE8g= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.4 h1:CXV68E2dNqhuynZJPB80bhPQwAKqBWVer887figW6Jc= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.4/go.mod h1:/xFi9KtvBXP97ppCz1TAEvU1Uf66qvid89rbem3wCzQ= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.18 h1:vvbXsA2TVO80/KT7ZqCbx934dt6PY+vQ8hZpUZ/cpYg= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.18/go.mod h1:m2JJHledjBGNMsLOF1g9gbAxprzq3KjC8e4lxtn+eWg= github.com/aws/aws-sdk-go-v2/service/sso v1.25.6 h1:rGtWqkQbPk7Bkwuv3NzpE/scwwL9sC1Ul3tn9x83DUI= github.com/aws/aws-sdk-go-v2/service/sso v1.25.6/go.mod h1:u4ku9OLv4TO4bCPdxf4fA1upaMaJmP9ZijGk3AAOC6Q= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4 h1:OV/pxyXh+eMA0TExHEC4jyWdumLxNbzz1P0zJoezkJc= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4/go.mod h1:8Mm5VGYwtm+r305FfPSuc+aFkrypeylGYhFim6XEPoc= github.com/aws/aws-sdk-go-v2/service/sts v1.34.1 h1:aUrLQwJfZtwv3/ZNG2xRtEen+NqI3iesuacjP51Mv1s= github.com/aws/aws-sdk-go-v2/service/sts v1.34.1/go.mod h1:3wFBZKoWnX3r+Sm7in79i54fBmNfwhdNdQuscCw7QIk= github.com/aws/smithy-go v1.22.4 h1:uqXzVZNuNexwc/xrh6Tb56u89WDlJY6HS+KC0S4QSjw= github.com/aws/smithy-go v1.22.4/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8= github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA= github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/charmbracelet/bubbles v0.21.0 h1:9TdC97SdRVg/1aaXNVWfFH3nnLAwOXr8Fn6u6mfQdFs= github.com/charmbracelet/bubbles v0.21.0/go.mod h1:HF+v6QUR4HkEpz62dx7ym2xc71/KBHg+zKwJtMw+qtg= github.com/charmbracelet/bubbletea v1.3.5 h1:JAMNLTbqMOhSwoELIr0qyP4VidFq72/6E9j7HHmRKQc= github.com/charmbracelet/bubbletea v1.3.5/go.mod h1:TkCnmH+aBd4LrXhXcqrKiYwRs7qyQx5rBgH5fVY3v54= github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs= github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk= github.com/charmbracelet/glamour v0.10.0 h1:MtZvfwsYCx8jEPFJm3rIBFIMZUfUJ765oX8V6kXldcY= github.com/charmbracelet/glamour v0.10.0/go.mod h1:f+uf+I/ChNmqo087elLnVdCiVgjSKWuXa/l6NU2ndYk= github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE= github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA= github.com/charmbracelet/x/ansi v0.8.0 h1:9GTq3xq9caJW8ZrBTe0LIe2fvfLR/bYXKTx2llXn7xE= github.com/charmbracelet/x/ansi v0.8.0/go.mod h1:wdYl/ONOLHLIVmQaxbIYEC/cRKOQyjTkowiI4blgS9Q= github.com/charmbracelet/x/cellbuf v0.0.13 h1:/KBBKHuVRbq1lYx5BzEHBAFBP8VcQzJejZ/IA3iR28k= github.com/charmbracelet/x/cellbuf v0.0.13/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs= github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91 h1:payRxjMjKgx2PaCWLZ4p3ro9y97+TVLZNaRZgJwSVDQ= github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U= github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf h1:rLG0Yb6MQSDKdB52aGX55JT1oi0P0Kuaj7wi1bLUpnI= github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf/go.mod h1:B3UgsnsBZS/eX42BlaNiJkD1pPOUa+oF1IYC6Yd2CEU= github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= github.com/chzyer/logex v1.2.1 h1:XHDu3E6q+gdHgsdTPH6ImJMIp436vR6MPtH8gP05QzM= github.com/chzyer/logex v1.2.1/go.mod h1:JLbx6lG2kDbNRFnfkgvh4eRJRPX1QCoOIWomwysCBrQ= github.com/chzyer/readline v1.5.1 h1:upd/6fQk4src78LMRzh5vItIt361/o4uq553V8B5sGI= github.com/chzyer/readline v1.5.1/go.mod h1:Eh+b79XXUwfKfcPLepksvw2tcLE/Ct21YObkaSkeBlk= github.com/chzyer/test v1.0.0 h1:p3BQDXSxOhOG0P9z6/hGnII4LGiEPOYBhs8asl/fC04= github.com/chzyer/test v1.0.0/go.mod h1:2JlltgoNkt4TW/z9V/IzDdFaMTM2JPIi26O1pF38GC8= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dlclark/regexp2 v1.11.4 h1:rPYF9/LECdNymJufQKmri9gV604RvvABwgOA8un7yAo= github.com/dlclark/regexp2 v1.11.4/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI= github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ= github.com/emicklei/go-restful/v3 v3.12.2 h1:DhwDP0vY3k8ZzE0RunuJy8GhNpPL6zqLkDf9B/a0/xU= github.com/emicklei/go-restful/v3 v3.12.2/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fxamacker/cbor/v2 v2.9.0 h1:NpKPmjDBgUfBms6tr6JZkTHtfFGcMKsw3eGcmD/sapM= github.com/fxamacker/cbor/v2 v2.9.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= github.com/go-openapi/jsonpointer v0.21.0 h1:YgdVicSA9vH5RiHs9TZW5oyafXZFc6+2Vc1rr/O9oNQ= github.com/go-openapi/jsonpointer v0.21.0/go.mod h1:IUyH9l/+uyhIYQ/PXVA41Rexl+kOkAPDdXEYns6fzUY= github.com/go-openapi/jsonreference v0.20.2 h1:3sVjiK66+uXK/6oQ8xgcRKcFgQ5KXa2KvnJRumpMGbE= github.com/go-openapi/jsonreference v0.20.2/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En5Ap4rVB5KVcIDZG2k= github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+GrE= github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ= github.com/go-quicktest/qt v1.101.0 h1:O1K29Txy5P2OK0dGo59b7b0LR6wKfIhttaAhHUyn7eI= github.com/go-quicktest/qt v1.101.0/go.mod h1:14Bz/f7NwaXPtdYEgzsx46kqSxVwTbzVZsDC26tQJow= github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI= github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/gnostic-models v0.7.0 h1:qwTtogB15McXDaNqTZdzPJRHvaVJlAl+HVQnLmJEJxo= github.com/google/gnostic-models v0.7.0/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20241029153458-d1b30febd7db h1:097atOisP2aRj7vFgYQBbFN4U4JNXUNYpxael3UzMyo= github.com/google/pprof v0.0.0-20241029153458-d1b30febd7db/go.mod h1:vavhavw2zAxS5dIdcRluK6cSGGPlZynqzFM8NdvU144= github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw= github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA= github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q= github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA= github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 h1:JeSE6pjso5THxAzdVpqr6/geYxZytqFMBCOtn/ujyeo= github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674/go.mod h1:r4w70xmWCQKmi1ONH4KIaBptdivuRPyosB9RmPlGEwA= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/keybase/go-keychain v0.0.1 h1:way+bWYa6lDppZoZcgMbYsvC7GxljxrskdNInRtuthU= github.com/keybase/go-keychain v0.0.1/go.mod h1:PdEILRW3i9D8JcdM+FmY6RwkHGnhHxXwkPPMeUgOK1k= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/mark3labs/mcp-go v0.41.1 h1:w78eWfiQam2i8ICL7AL0WFiq7KHNJQ6UB53ZVtH4KGA= github.com/mark3labs/mcp-go v0.41.1/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= github.com/moby/spdystream v0.5.0 h1:7r0J1Si3QO/kjRitvSLVVFUjxMEb/YLj6S9FF62JBCU= github.com/moby/spdystream v0.5.0/go.mod h1:xBAYlnt/ay+11ShkdFKNAG7LsyK/tmNBVvVOwrfMgdI= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee h1:W5t00kpgFdJifH4BDsTlE89Zl93FEloxaWZfGcifgq8= github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus= github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f/go.mod h1:ZdcZmHo+o7JKHSa8/e818NopupXU1YMK5fe1lsApnBw= github.com/ollama/ollama v0.6.5 h1:vXKkVX57ql/1ZzMw4SVK866Qfd6pjwEcITVyEpF0QXQ= github.com/ollama/ollama v0.6.5/go.mod h1:pGgtoNyc9DdM6oZI6yMfI6jTk2Eh4c36c2GpfQCH7PY= github.com/onsi/ginkgo/v2 v2.21.0 h1:7rg/4f3rB88pb5obDgNZrNHrQ4e6WpjonchcpuBRnZM= github.com/onsi/ginkgo/v2 v2.21.0/go.mod h1:7Du3c42kxCUegi0IImZ1wUQzMBVecgIHjR1C+NkhLQo= github.com/onsi/gomega v1.35.1 h1:Cwbd75ZBPxFSuZ6T+rN/WCb/gOc6YgFBXLlZLhC7Ds4= github.com/onsi/gomega v1.35.1/go.mod h1:PvZbdDc8J6XJEpDK4HCuRBm8a6Fzp9/DmhC9C7yFlog= github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0= github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM= github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA= github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/sahilm/fuzzy v0.1.1 h1:ceu5RHF8DGgoi+/dR5PsECjCDH1BE3Fnmpo7aVXOdRA= github.com/sahilm/fuzzy v0.1.1/go.mod h1:VFvziUEIMCrT6A6tw2RFIXPXXmzXbOsSHF0DOI8ZK9Y= github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo= github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0= github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o= github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.7.1/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E= github.com/yuin/goldmark v1.7.8 h1:iERMLn0/QJeHFhxSt3p6PeN9mGnvIKSpG9YYorDMnic= github.com/yuin/goldmark v1.7.8/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E= github.com/yuin/goldmark-emoji v1.0.5 h1:EMVWyCGPlXJfUXBXpuMu+ii3TIaxbVBnEX9uaDC4cIk= github.com/yuin/goldmark-emoji v1.0.5/go.mod h1:tTkZEbwu5wkPmgTcitqddVxY9osFZiavD+r4AzQrh1U= go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 h1:CV7UdSGJt/Ao6Gp4CXckLxVRRsRgDHoI8XjbL3PDl8s= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0/go.mod h1:FRmFuRJfag1IZ2dPkHnEoSFVgTVPUd2qf5Vi69hLb8I= go.opentelemetry.io/otel v1.34.0 h1:zRLXxLCgL1WyKsPVrgbSdMN4c0FMkDAskSTQP+0hdUY= go.opentelemetry.io/otel v1.34.0/go.mod h1:OWFPOQ+h4G8xpyjgqo4SxJYdDQ/qmRH+wivy7zzx9oI= go.opentelemetry.io/otel/metric v1.34.0 h1:+eTR3U0MyfWjRDhmFMxe2SsW64QrZ84AOhvqS7Y+PoQ= go.opentelemetry.io/otel/metric v1.34.0/go.mod h1:CEDrp0fy2D0MvkXE+dPV7cMi8tWZwX3dmaIhwPOaqHE= go.opentelemetry.io/otel/sdk v1.34.0 h1:95zS4k/2GOy069d321O8jWgYsW3MzVV+KuSPKp7Wr1A= go.opentelemetry.io/otel/sdk v1.34.0/go.mod h1:0e/pNiaMAqaykJGKbi+tSjWfNNHMTxoC9qANsCzbyxU= go.opentelemetry.io/otel/sdk/metric v1.32.0 h1:rZvFnvmvawYb0alrYkjraqJq0Z4ZUJAiyYCU9snn1CU= go.opentelemetry.io/otel/sdk/metric v1.32.0/go.mod h1:PWeZlq0zt9YkYAp3gjKZ0eicRYvOh1Gd+X99x6GHpCQ= go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC8mh/k= go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI= go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/crypto v0.40.0 h1:r4x+VvoG5Fm+eJcxMaY8CQM7Lb0l1lsmjGBQ6s8BfKM= golang.org/x/crypto v0.40.0/go.mod h1:Qr1vMER5WyS2dfPHAlsOj01wgLbsyWtFn/aY+5+ZdxY= golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa h1:t2QcU6V556bFjYgu4L6C+6VrCPyJZ+eyRsABUPs1mz4= golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa/go.mod h1:BHOTPb3L19zxehTsLoJXVaTktb06DFgmdW6Wb9s8jqk= golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw= golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA= golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw= golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.34.0 h1:H5Y5sJ2L2JRdyv7ROF1he/lPdvFsd0mJHFw2ThKHxLA= golang.org/x/sys v0.34.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/term v0.33.0 h1:NuFncQrRcaRvVmgRkvM3j/F00gWIAlcmlB8ACEKmGIg= golang.org/x/term v0.33.0/go.mod h1:s18+ql9tYWp1IfpV9DmCtQDDSRBUjKaw9M1eAv5UeF0= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.27.0 h1:4fGWRpyh641NLlecmyl4LOe6yDdfaYNrGb2zdfo4JV4= golang.org/x/text v0.27.0/go.mod h1:1D28KMCvyooCX9hBiosv5Tz/+YLxj0j7XhWjpSUF7CU= golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4= golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg= golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/genai v1.8.0 h1:unX2CNWSiKDO2MSTKK3RstXg/vHp9hr42LIcL6f3Cik= google.golang.org/genai v1.8.0/go.mod h1:TyfOKRz/QyCaj6f/ZDt505x+YreXnY40l2I6k8TvgqY= google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2 h1:DMTIbak9GhdaSxEjvVzAeNZvyc03I61duqNbnm3SU0M= google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2/go.mod h1:LuRYeWDFV6WOn90g357N17oMCaxpgCnbi/44qJvDn2I= google.golang.org/grpc v1.70.0 h1:pWFv03aZoHzlRKHWicjsZytKAiYCtNS0dHbXnIdq7jQ= google.golang.org/grpc v1.70.0/go.mod h1:ofIJqVKDXx/JiXrwr2IG4/zwdH9txy3IlF40RmcJSQw= google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM= google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/evanphx/json-patch.v4 v4.12.0 h1:n6jtcsulIzXPJaxegRbvFNNrZDjbij7ny3gmSPG+6V4= gopkg.in/evanphx/json-patch.v4 v4.12.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M= gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc= gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw= gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/api v0.34.2 h1:fsSUNZhV+bnL6Aqrp6O7lMTy6o5x2C4XLjnh//8SLYY= k8s.io/api v0.34.2/go.mod h1:MMBPaWlED2a8w4RSeanD76f7opUoypY8TFYkSM+3XHw= k8s.io/apimachinery v0.34.2 h1:zQ12Uk3eMHPxrsbUJgNF8bTauTVR2WgqJsTmwTE/NW4= k8s.io/apimachinery v0.34.2/go.mod h1:/GwIlEcWuTX9zKIg2mbw0LRFIsXwrfoVxn+ef0X13lw= k8s.io/client-go v0.34.2 h1:Co6XiknN+uUZqiddlfAjT68184/37PS4QAzYvQvDR8M= k8s.io/client-go v0.34.2/go.mod h1:2VYDl1XXJsdcAxw7BenFslRQX28Dxz91U9MWKjX97fE= k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk= k8s.io/klog/v2 v2.130.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE= k8s.io/kube-openapi v0.0.0-20250710124328-f3f2b991d03b h1:MloQ9/bdJyIu9lb1PzujOPolHyvO06MXG5TUIj2mNAA= k8s.io/kube-openapi v0.0.0-20250710124328-f3f2b991d03b/go.mod h1:UZ2yyWbFTpuhSbFhv24aGNOdoRdJZgsIObGBUaYVsts= k8s.io/utils v0.0.0-20250604170112-4c0f3b243397 h1:hwvWFiBzdWw1FhfY1FooPn3kzWuJ8tmbZBHi4zVsl1Y= k8s.io/utils v0.0.0-20250604170112-4c0f3b243397/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0= mvdan.cc/sh/v3 v3.11.0 h1:q5h+XMDRfUGUedCqFFsjoFjrhwf2Mvtt1rkMvVz0blw= mvdan.cc/sh/v3 v3.11.0/go.mod h1:LRM+1NjoYCzuq/WZ6y44x14YNAI0NK7FLPeQSaFagGg= sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8 h1:gBQPwqORJ8d8/YNZWEjoZs7npUVDpVXUUOFfW6CgAqE= sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8/go.mod h1:mdzfpAEoE6DHQEN0uh9ZbOCuHbLK5wOm7dK4ctXE9Tg= sigs.k8s.io/randfill v1.0.0 h1:JfjMILfT8A6RbawdsK2JXGBR5AQVfd+9TbzrlneTyrU= sigs.k8s.io/randfill v1.0.0/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= sigs.k8s.io/structured-merge-diff/v6 v6.3.0 h1:jTijUJbW353oVOd9oTlifJqOGEkUw2jB/fXCbTiQEco= sigs.k8s.io/structured-merge-diff/v6 v6.3.0/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= sigs.k8s.io/yaml v1.6.0/go.mod h1:796bPqUfzR/0jLAl6XjHl3Ck7MiyVv8dbTdyT3/pMf4= ================================================ FILE: go.work ================================================ go 1.24.0 toolchain go1.24.4 use ( . ./gollm ./kubectl-utils ) ================================================ FILE: go.work.sum ================================================ cel.dev/expr v0.19.1/go.mod h1:MrpN08Q+lEBs+bGYdLxxHkZoUSsCp0nSKTs0nTymJgw= cloud.google.com/go/auth/oauth2adapt v0.2.7/go.mod h1:NTbTTzfvPl1Y3V1nPpOgl2w6d/FjO7NNUQaWSox6ZMc= cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k= cloud.google.com/go/iam v1.3.1/go.mod h1:3wMtuyT4NcbnYNPLMBzYRFiEfjKfJlLVLrisE7bwm34= cloud.google.com/go/longrunning v0.5.6/go.mod h1:vUaDrWYOMKRuhiv6JBnn49YxCPz2Ayn9GqyjaBT8/mA= cloud.google.com/go/monitoring v1.23.0/go.mod h1:034NnlQPDzrQ64G2Gavhl0LUHZs9H3rRmhtnp7jiJgg= cloud.google.com/go/storage v1.50.0/go.mod h1:l7XeiD//vx5lfqE3RavfmU9yvk5Pp0Zhcv482poyafY= cloud.google.com/go/translate v1.10.3/go.mod h1:GW0vC1qvPtd3pgtypCv4k4U8B7EdgK9/QEF2aJEUovs= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0/go.mod h1:XCW7KnZet0Opnr7HccfUw1PLc4CjHqpcaxW8DHklNkQ= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.1/go.mod h1:zGqV2R4Cr/k8Uye5w+dgQ06WJtEcbQG/8J7BB6hnCr4= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0/go.mod h1:9kIvujWAA58nmPmWB1m23fyWic1kYZMxD9CxaWn4Qpg= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.2/go.mod h1:SqINnQ9lVVdRlyC8cd1lCI0SdX4n2paeABd2K8ggfnE= github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0/go.mod h1:iZDifYGJTIgIIkYRNWPENUnqx6bJ2xnSDFI2tjwZNuY= github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.0/go.mod h1:j2chePtV91HrC22tGoRX3sGY42uF13WzmmV80/OdVAA= github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI= github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.25.0/go.mod h1:obipzmGjfSjam60XLwGfqUkJsfiheAl+TUjG+4yzyPM= github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.49.0/go.mod h1:6fTWu4m3jocfUZLYF5KsZC1TUfRvEjs7lM4crme/irw= github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.49.0/go.mod h1:wRbFgBQUVm1YXrvWKofAEmq9HNJTDphbAaJSSX01KUI= github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= github.com/NYTimes/gziphandler v1.1.1/go.mod h1:n/CVRwUEOgIxrgPvAQhUUr9oeUtvrhMomdKFjzJNB0c= github.com/agnivade/levenshtein v1.1.1/go.mod h1:veldBMzWxcCG2ZvUTKD2kJNRdCk5hVbJomOvKkmgYbo= github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40/go.mod h1:Q7yQnSMnLvcXlZ8RV+jwz/6y1rQTqbX6C82SndT52Zs= github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4= github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= github.com/census-instrumentation/opencensus-proto v0.4.1/go.mod h1:4T9NM4+4Vw91VeyqjLS6ao50K5bOcLKN6Q42XnYaRYw= github.com/charmbracelet/harmonica v0.2.0/go.mod h1:KSri/1RMQOZLbw7AHqgcBycp8pgJnQMYYT8QZRqZ1Ao= github.com/chewxy/hm v1.0.0/go.mod h1:qg9YI4q6Fkj/whwHR1D+bOGeF7SniIP40VweVepLjg0= github.com/chewxy/math32 v1.11.0/go.mod h1:dOB2rcuFrCn6UHrze36WSLVPKtzPMRAQvBvUwkSsLqs= github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= github.com/cncf/xds/go v0.0.0-20241223141626-cff3c89139a3/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8= github.com/containerd/console v1.0.3/go.mod h1:7LqA/THxQ86k76b8c/EMSiaJ3h1eZkMkXar0TQ1gf3U= github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1/go.mod h1:uw2gLcxEuYUlAd/EXyjc/v55nd3+47YAgWbSXVxPrNI= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/emirpasic/gods/v2 v2.0.0-alpha/go.mod h1:W0y4M2dtBB9U5z3YlghmpuUhiaZT2h6yoeE+C1sCp6A= github.com/envoyproxy/go-control-plane v0.13.4/go.mod h1:kDfuBlDVsSj2MjrLEtRWtHlsWIFcGyB2RMO44Dc5GZA= github.com/envoyproxy/go-control-plane/envoy v1.32.4/go.mod h1:Gzjc5k8JcJswLjAx1Zm+wSYE20UrLtt7JZMWiWQXQEw= github.com/envoyproxy/go-control-plane/ratelimit v0.1.0/go.mod h1:Wk+tMFAFbCXaJPzVVHnPgRKdUdwW/KdbRt94AzgRee4= github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU= github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= github.com/gin-contrib/cors v1.7.2/go.mod h1:SUJVARKgQ40dmrzgXEVxj2m7Ig1v1qIboQkPDTQ9t2E= github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= github.com/gin-gonic/gin v1.10.0/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y= github.com/go-logr/logr v1.2.0/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.1/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-openapi/jsonreference v0.20.1/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En5Ap4rVB5KVcIDZG2k= github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM= github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= github.com/golang-jwt/jwt/v4 v4.4.3/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= github.com/golang-jwt/jwt/v5 v5.0.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= github.com/golang/glog v1.2.4/go.mod h1:6AhwSGph0fcJtXVM/PEHPqZlFeoLxhs7/t5UDAwmO+w= github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4= github.com/google/flatbuffers v24.3.25+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-pkcs11 v0.3.0/go.mod h1:6eQoGcuNJpa7jnd5pMGdkSaQpNDYvPlXWMcjXXThLlY= github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0= github.com/google/pprof v0.0.0-20240727154555-813a5fbdbec8/go.mod h1:K1liHPHnj73Fdn/EKuT8nrFqBihUSKXoLYU0BuatOYo= github.com/google/pprof v0.0.0-20240827171923-fa2c70bbbfe5/go.mod h1:vavhavw2zAxS5dIdcRluK6cSGGPlZynqzFM8NdvU144= github.com/google/renameio/v2 v2.0.0/go.mod h1:BtmJXm5YlszgC+TD4HOEEUFgkJP3nLxehU6hfe7jRt4= github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA= github.com/iancoleman/strcase v0.3.0/go.mod h1:iwCmte+B7n89clKwxIoIXy/HfoL7AsD47ZCWhYzw7ho= github.com/ianlancetaylor/demangle v0.0.0-20240312041847-bd984b5ce465/go.mod h1:gx7rwoVhcfuVKG5uya9Hs3Sxj7EIvldVofAWIUtGouw= github.com/keybase/dbus v0.0.0-20220506165403-5aa21ea2c23a/go.mod h1:YPNKjjE7Ubp9dTbnWvsP3HT+hYnY6TfXzubYTBeUxc8= github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/lyft/protoc-gen-star/v2 v2.0.4-0.20230330145011-496ad1ac90a4/go.mod h1:amey7yeodaJhXSbf/TlLvWiqQfLOSpEk//mLlc+axEk= github.com/montanaflynn/stats v0.7.0/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow= github.com/nlpodyssey/gopickle v0.3.0/go.mod h1:f070HJ/yR+eLi5WmM1OXJEGaTpuJEUiib19olXgYha0= github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= github.com/onsi/ginkgo/v2 v2.19.0/go.mod h1:rlwLi9PilAFJ8jCg9UE1QP6VBpd6/xj3SRC0d6TU0To= github.com/onsi/ginkgo/v2 v2.20.1/go.mod h1:lG9ey2Z29hR41WMVthyJBGUBcBhGOtoPF2VFMvBXFCI= github.com/onsi/gomega v1.33.1/go.mod h1:U4R44UsT+9eLIaYRB2a5qajjtQYn0hauxvRm16AVYg0= github.com/onsi/gomega v1.34.2/go.mod h1:v1xfxRgk0KIsG+QOdm7p8UosrOzPYRo60fd3B/1Dukc= github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c/go.mod h1:PSojXDXF7TbgQiD6kkd98IHOS0QqTyUEaWRiS8+BLu8= github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU= github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8/go.mod h1:HKlIX3XHQyzLZPlr7++PzdhaXEj94dEiJgZDTsxEqUI= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/spf13/afero v1.10.0/go.mod h1:UBogFpq8E9Hx+xc5CNTTEpTnuHVmXDwZcZcE1eb/UhQ= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= github.com/xtgo/set v1.0.0/go.mod h1:d3NHzGzSa0NmB2NhFyECA+QdRp29oEn2xbT+TpeFoM8= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo= go.opentelemetry.io/contrib/detectors/gcp v1.34.0/go.mod h1:cV4BMFcscUR/ckqLkbfQmF0PRsq8w/lMGzdbCSveBHo= go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.59.0/go.mod h1:ijPqXp5P6IRRByFVVg9DY8P5HkxkHE5ARIa+86aXPf4= go.opentelemetry.io/otel/sdk/metric v1.34.0/go.mod h1:jQ/r8Ze28zRKoNRdkjCZxfs6YvBTG1+YIqyFVFYec5w= go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.3/go.mod h1:tBHosrYAkRZjRAOREWbDnBXUf08JOwYq++0QNwQiWzI= go4.org/unsafe/assume-no-moving-gc v0.0.0-20231121144256-b99613f794b6/go.mod h1:FftLjUGFEDu5k8lt0ddY+HcrH/qU/0qk+H8j9/nTl3E= golang.org/x/arch v0.8.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc= golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc= golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U= golang.org/x/crypto v0.41.0 h1:WKYxWedPGCTVVl5+WHSSrOBT0O8lx32+zxmHxijgXp4= golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sUc= golang.org/x/image v0.22.0/go.mod h1:9hPFhljd4zZ1GNSIZJ49sqbp45GKK9t6w+iXvGqZUz4= golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/mod v0.20.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/mod v0.21.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY= golang.org/x/mod v0.23.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY= golang.org/x/mod v0.25.0/go.mod h1:IXM97Txy2VM4PJ3gI61r1YEk/gAj6zAHN3AdZt6S9Ww= golang.org/x/mod v0.26.0/go.mod h1:/j6NAhSk8iQ723BGAUyoAcn7SlD7s15Dp9Nd/SfeaFQ= golang.org/x/mod v0.27.0/go.mod h1:rWI627Fq0DEoudcK+MBkNkCe0EetEaDSwJJkCcjpazc= golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44= golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE= golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= golang.org/x/net v0.30.0/go.mod h1:2wGyMJ5iFasEhkwi13ChkO/t1ECNC4X4eBKkVFyYFlU= golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k= golang.org/x/net v0.37.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= golang.org/x/net v0.42.0/go.mod h1:FF1RA5d3u7nAYA4z2TkclSCKh68eSXtiFwcWQpPXdt8= golang.org/x/net v0.43.0 h1:lat02VYK2j4aLzMzecihNvTlJNQUq316m2Mr9rnM6YE= golang.org/x/net v0.43.0/go.mod h1:vhO1fvI4dGsIjh73sWfUVjj3N7CA9WkKJNQm2svM6Jg= golang.org/x/oauth2 v0.25.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI= golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20210616045830-e2b7044e8c71/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.23.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/telemetry v0.0.0-20250807160809-1a19826ec488/go.mod h1:fGb/2+tgXXjhjHsTNdVEEMZNWA0quBnfrO+AfoDSAKw= golang.org/x/term v0.32.0/go.mod h1:uZG1FhGx848Sqfsq4/DlJr3xGGsYMu/L5GW4abiaEPQ= golang.org/x/term v0.34.0 h1:O/2T7POpk0ZZ7MAzMeWFSg6S5IpWd/RXDlM9hgM3DR4= golang.org/x/term v0.34.0/go.mod h1:5jC53AEywhIVebHgPVeg0mj8OD3VO9OzclacVrqpaAw= golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= golang.org/x/text v0.19.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA= golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng= golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU= golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk= golang.org/x/tools v0.24.0/go.mod h1:YhNqVBIfWHdzvTLs0d8LCuMhkKUgSUKldakyV7W/WDQ= golang.org/x/tools v0.30.0/go.mod h1:c347cR/OJfw5TI+GfX7RUPNMdDRRbjvYTS0jPyvsVtY= golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= golang.org/x/tools v0.35.0/go.mod h1:NKdj5HkL/73byiZSJjqJgKn3ep7KjFkBOkR/Hps3VPw= gonum.org/v1/gonum v0.15.0/go.mod h1:xzZVBJBtS+Mz4q0Yl2LJTk+OxOg4jiXZ7qBoM0uISGo= google.golang.org/api v0.222.0/go.mod h1:efZia3nXpWELrwMlN5vyQrD4GmJN1Vw0x68Et3r+a9c= google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds= google.golang.org/genproto v0.0.0-20250122153221-138b5a5a4fd4 h1:Pw6WnI9W/LIdRxqK7T6XGugGbHIRl5Q7q3BssH6xk4s= google.golang.org/genproto v0.0.0-20250122153221-138b5a5a4fd4/go.mod h1:qbZzneIOXSq+KFAFut9krLfRLZiFLzZL5u2t8SV83EE= google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422/go.mod h1:b6h1vNKhxaSoEI+5jc3PJUCustfli/mRab7295pY7rw= google.golang.org/genproto/googleapis/rpc v0.0.0-20240826202546-f6391c0de4c7/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU= google.golang.org/genproto/googleapis/rpc v0.0.0-20250115164207-1a7da9e5054f/go.mod h1:+2Yz8+CLJbIfL9z73EW45avw8Lmge3xVElCP9zEKi50= google.golang.org/grpc v1.71.0 h1:kF77BGdPTQ4/JZWMlb9VpJ5pa25aqvVqogsxNHHdeBg= google.golang.org/grpc v1.71.0/go.mod h1:H0GRtasmQOh9LkFoCPDu3ZrwUtD1YGE+b2vYBYd/8Ec= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= google.golang.org/protobuf v1.35.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= google.golang.org/protobuf v1.36.4/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= gorgonia.org/vecf32 v0.9.0/go.mod h1:NCc+5D2oxddRL11hd+pCB1PEyXWOyiQxfZ/1wwhOXCA= gorgonia.org/vecf64 v0.9.0/go.mod h1:hp7IOWCnRiVQKON73kkC/AUMtEXyf9kGlVrtPQ9ccVA= k8s.io/gengo/v2 v2.0.0-20240826214909-a7b603a56eb7/go.mod h1:EJykeLsmFC60UQbYJezXkEsG2FLrt0GPNkU5iK5GWxU= k8s.io/gengo/v2 v2.0.0-20250604051438-85fd79dbfd9f/go.mod h1:EJykeLsmFC60UQbYJezXkEsG2FLrt0GPNkU5iK5GWxU= k8s.io/klog/v2 v2.80.1/go.mod h1:y1WjHnz7Dj687irZUWR/WLkLc5N1YHtjLdmgWjndZn0= k8s.io/utils v0.0.0-20240711033017-18e509b52bc8/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0= mvdan.cc/editorconfig v0.3.0/go.mod h1:NcJHuDtNOTEJ6251indKiWuzK6+VcrMuLzGMLKBFupQ= sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd/go.mod h1:B8JuhiUyNFVKdsE8h686QcCxMaH6HrOAZj4vswFpcB0= sigs.k8s.io/structured-merge-diff/v6 v6.2.0/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= ================================================ FILE: gollm/README.md ================================================ # gollm A Go library for calling into multiple Large Language Model (LLM) providers with a unified interface. This library is intended for use by kubectl-ai, but may prove useful for other similar go tools in future. Note that the library is still evolving and will likely make incompatible changes often. We are focusing on kubectl-ai's use-case, but will consider changes to support additional use-cases. ## Overview gollm provides a consistent API for interacting with various LLM providers, making it easy to switch between different models and services without changing your application code. The library supports both chat-based conversations and single completions, with features like function calling, streaming responses, and retry logic. ## Features - **Multi-provider support**: OpenAI, Azure OpenAI, Google Gemini, Ollama, LlamaCPP, Grok, and more - **Unified interface**: Consistent API across all providers - **Chat conversations**: Multi-turn conversations with conversation history - **Function calling**: Define and use custom functions with LLMs - **Streaming support**: Real-time streaming responses - **Retry logic**: Built-in retry mechanisms with configurable backoff - **Response schemas**: Constrain LLM responses to specific JSON schemas - **SSL configuration**: Optional SSL certificate verification skipping - **Environment-based configuration**: Easy setup via environment variables ## Providers | Provider | ID | Description | |----------|----|-------------| | OpenAI | `openai://` | OpenAI's GPT models | | Azure OpenAI | `azopenai://` | Microsoft Azure's OpenAI service | | Google Gemini | `gemini://` | Google's Gemini models | | Vertex AI | `vertexai://` | Google Cloud Vertex AI (via Gemini) | | Ollama | `ollama://` | Local Ollama models | | LlamaCPP | `llamacpp://` | Local LlamaCPP models | | Grok | `grok://` | xAI's Grok models | | Anthropic | `anthropic://` | Claude models with native tool use, prompt caching, and extended thinking | ## Quick Start ### Installation ```bash go get github.com/GoogleCloudPlatform/kubectl-ai/gollm ``` ### Basic Usage ```go package main import ( "context" "fmt" "log" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" ) func main() { ctx := context.Background() // Create a client using environment variable client, err := gollm.NewClient(ctx, "") if err != nil { log.Fatal(err) } defer client.Close() // Start a chat conversation chat := client.StartChat("You are a helpful assistant.", "gpt-3.5-turbo") // Send a message response, err := chat.Send(ctx, "Hello, how are you?") if err != nil { log.Fatal(err) } // Print the response for _, candidate := range response.Candidates() { fmt.Println(candidate.String()) } } ``` ### Environment Configuration Set the `LLM_CLIENT` environment variable to specify your preferred provider: ```bash # OpenAI export LLM_CLIENT="openai://api.openai.com" export OPENAI_API_KEY="your-api-key" # Azure OpenAI export LLM_CLIENT="azopenai://your-resource.openai.azure.com" export AZURE_OPENAI_API_KEY="your-api-key" # Google Gemini export LLM_CLIENT="gemini://generativelanguage.googleapis.com" export GOOGLE_API_KEY="your-api-key" # Ollama (local) export LLM_CLIENT="ollama://localhost:11434" ``` ## Examples ### Single Completion ```go ctx := context.Background() client, err := gollm.NewClient(ctx, "openai://api.openai.com") if err != nil { log.Fatal(err) } defer client.Close() req := &gollm.CompletionRequest{ Model: "gpt-3.5-turbo", Prompt: "Write a short poem about programming", } response, err := client.GenerateCompletion(ctx, req) if err != nil { log.Fatal(err) } fmt.Println(response.Response()) ``` ### Streaming Chat ```go ctx := context.Background() client, err := gollm.NewClient(ctx, "openai://api.openai.com") if err != nil { log.Fatal(err) } defer client.Close() chat := client.StartChat("You are a helpful assistant.", "gpt-3.5-turbo") // Send a streaming message iterator, err := chat.SendStreaming(ctx, "Tell me a story about a robot") if err != nil { log.Fatal(err) } // Process streaming response for response := range iterator { if response.V1 != nil { for _, candidate := range response.V1.Candidates() { for _, part := range candidate.Parts() { if text, ok := part.AsText(); ok { fmt.Print(text) } } } } if response.V2 != nil { // Handle error log.Printf("Error: %v", response.V2) break } } ``` ### Function Calling ```go // Define a function that the LLM can call functionDef := &gollm.FunctionDefinition{ Name: "get_weather", Description: "Get the current weather for a location", Parameters: &gollm.Schema{ Type: gollm.TypeObject, Properties: map[string]*gollm.Schema{ "location": { Type: gollm.TypeString, Description: "The city and state, e.g. San Francisco, CA", }, "unit": { Type: gollm.TypeString, Description: "The temperature unit to use. Infer this from the user's location.", Required: []string{"location"}, }, }, }, } chat := client.StartChat("You are a helpful assistant.", "gpt-3.5-turbo") chat.SetFunctionDefinitions([]*gollm.FunctionDefinition{functionDef}) response, err := chat.Send(ctx, "What's the weather like in San Francisco?") if err != nil { log.Fatal(err) } // Check for function calls in the response for _, candidate := range response.Candidates() { for _, part := range candidate.Parts() { if functionCalls, ok := part.AsFunctionCalls(); ok { for _, call := range functionCalls { fmt.Printf("Function call: %s with args %v\n", call.Name, call.Arguments) // Execute the function and send the result back result := executeWeatherFunction(call.Arguments) chat.Send(ctx, gollm.FunctionCallResult{ ID: call.ID, Name: call.Name, Result: result, }) } } } } ``` ### Response Schema Constraints ```go // Define a schema for structured responses schema := &gollm.Schema{ Type: gollm.TypeObject, Properties: map[string]*gollm.Schema{ "name": { Type: gollm.TypeString, Description: "The person's name", }, "age": { Type: gollm.TypeInteger, Description: "The person's age", }, "interests": { Type: gollm.TypeArray, Items: &gollm.Schema{ Type: gollm.TypeString, }, Description: "List of interests", }, }, Required: []string{"name", "age"}, } client.SetResponseSchema(schema) // Now all responses will be constrained to match this schema response, err := chat.Send(ctx, "Tell me about a person named Alice who is 30 years old") ``` ### Retry Logic ```go // Configure retry behavior retryConfig := gollm.RetryConfig{ MaxAttempts: 3, InitialBackoff: time.Second, MaxBackoff: 30 * time.Second, BackoffFactor: 2.0, Jitter: true, } // Create a chat with retry logic chat := client.StartChat("You are a helpful assistant.", "gpt-3.5-turbo") retryChat := gollm.NewRetryChat(chat, retryConfig) // Use the retry chat - it will automatically retry on retryable errors response, err := retryChat.Send(ctx, "Hello!") ``` ### Building Schemas from Go Types ```go type Person struct { Name string `json:"name"` Age int `json:"age"` Interests []string `json:"interests,omitempty"` } // Automatically build a schema from a Go struct schema := gollm.BuildSchemaFor(reflect.TypeOf(Person{})) // Use the schema to constrain responses client.SetResponseSchema(schema) ``` ## Configuration Options ### Client Options ```go // Create a client with custom options client, err := gollm.NewClient(ctx, "openai://api.openai.com", gollm.WithSkipVerifySSL(), // Skip SSL verification (for development) ) ``` ### Environment Variables - `LLM_CLIENT`: The provider URL to use (e.g., "openai://api.openai.com") - `LLM_SKIP_VERIFY_SSL`: Set to "1" or "true" to skip SSL certificate verification - Provider-specific API keys (e.g., `OPENAI_API_KEY`, `GOOGLE_API_KEY`)

Anthropic-specific environment variables and provider features

 #### Anthropic-specific environment variables | Variable | Description | Default | |----------|-------------|---------| | `ANTHROPIC_API_KEY` | Anthropic API key (required) | — | | `ANTHROPIC_MODEL` | Default Claude model to use | `claude-sonnet-4-6` | | `ANTHROPIC_PROMPT_CACHING` | Enable prompt caching (`"false"` to disable) | `true` | | `ANTHROPIC_EXTENDED_THINKING` | Enable extended thinking(`"true"` to enable) | `false` | | `ANTHROPIC_MAX_TOKENS` | Max output tokens per request | `4096` | ### Anthropic provider features #### Prompt caching Enabled by default. Anthropic's [prompt caching](https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching) attaches a `cache_control` breakpoint to the system prompt and the last tool definition so that both are reused from cache on every subsequent turn. Because kubectl-ai's system prompt is large and repeated verbatim each turn, this typically cuts input token costs significantly after the first request. Set `ANTHROPIC_PROMPT_CACHING=false` to disable. #### Extended thinking Disabled by default. When enabled, Claude produces a `thinking` content block containing its internal reasoning before giving a final answer. This can improve accuracy on complex multi-step queries (e.g. root-cause analysis across multiple Kubernetes resources). Requires a model that supports extended thinking (`claude-3-7-sonnet-20250219` or later) and reserves 8 000 tokens for the thinking budget. The thinking block is kept in the conversation history (required by the API for multi-turn consistency) but is **not** shown in the terminal output. ```bash ANTHROPIC_EXTENDED_THINKING=true \ kubectl-ai --llm-provider=anthropic --model claude-3-7-sonnet-20250219 \ "why is my pod crashlooping" ``` Set `ANTHROPIC_EXTENDED_THINKING=true` to enable. #### Native streaming The Anthropic provider uses the official SSE event stream directly, bypassing the OpenAI compatibility shim. Tool input JSON is accumulated across `content_block_delta` events and only emitted as a complete `FunctionCall` once the block closes, so partial JSON is never forwarded to the agent loop. #### Retryable errors The provider maps Anthropic-native HTTP status codes to retry decisions: | Status | Meaning | Retried? | |--------|---------|----------| | 429 | Rate limit | Yes | | 529 | Overloaded | Yes | | 5xx | Server error | Yes | | 4xx (other) | Client error | No |
## Error Handling The library provides structured error handling with retryable error detection: ```go var apiErr *gollm.APIError if errors.As(err, &apiErr) { fmt.Printf("API Error: Status=%d, Message=%s\n", apiErr.StatusCode, apiErr.Message) } // Check if an error is retryable if chat.IsRetryableError(err) { // Implement retry logic } ``` ## Adding a provider To add a new provider: 1. Create a new file (e.g., `myprovider.go`) 2. Implement the `Client` interface 3. Register the provider in an `init()` function: ```go func init() { if err := gollm.RegisterProvider("myprovider", myProviderFactory); err != nil { panic(err) } } ``` ## License This project is licensed under the Apache License, Version 2.0. See the LICENSE file for details. ================================================ FILE: gollm/anthropic.go ================================================ // Copyright 2026 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "errors" "fmt" "os" "strconv" "strings" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" anthropic "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/option" "k8s.io/klog/v2" ) // Package-level env var storage (Anthropic env) var ( anthropicAPIKey string anthropicDefaultModel string anthropicPromptCaching bool anthropicExtendedThinking bool anthropicMaxTokens int64 ) func init() { anthropicAPIKey = os.Getenv("ANTHROPIC_API_KEY") anthropicDefaultModel = os.Getenv("ANTHROPIC_MODEL") // Prompt caching defaults to true; set ANTHROPIC_PROMPT_CACHING=false to disable if v := os.Getenv("ANTHROPIC_PROMPT_CACHING"); v == "false" { anthropicPromptCaching = false } else { anthropicPromptCaching = true } if v := os.Getenv("ANTHROPIC_EXTENDED_THINKING"); strings.ToLower(v) == "true" { anthropicExtendedThinking = true } anthropicMaxTokens = 4096 // default if v := os.Getenv("ANTHROPIC_MAX_TOKENS"); v != "" { if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { anthropicMaxTokens = n } } if err := RegisterProvider("anthropic", newAnthropicClientFactory); err != nil { klog.Fatalf("Failed to register anthropic provider: %v", err) } } // newAnthropicClientFactory creates a new Anthropic client with the given options. func newAnthropicClientFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewAnthropicClient(ctx, opts) } // AnthropicClient implements the gollm.Client interface for Anthropic models. type AnthropicClient struct { client *anthropic.Client } // Ensure AnthropicClient implements the Client interface. var _ Client = &AnthropicClient{} // NewAnthropicClient creates a new client for interacting with Anthropic models. func NewAnthropicClient(ctx context.Context, opts ClientOptions) (*AnthropicClient, error) { apiKey := anthropicAPIKey if apiKey == "" { return nil, errors.New("Anthropic API key not found. Set via ANTHROPIC_API_KEY env var") } httpClient := createCustomHTTPClient(opts.SkipVerifySSL) httpClient = withJournaling(httpClient) clientOpts := []option.RequestOption{ option.WithAPIKey(apiKey), option.WithHTTPClient(httpClient), } client := anthropic.NewClient(clientOpts...) return &AnthropicClient{client: &client}, nil } // Close cleans up any resources used by the client. func (c *AnthropicClient) Close() error { return nil } // StartChat starts a new chat session with the specified system prompt and model. func (c *AnthropicClient) StartChat(systemPrompt, model string) Chat { selectedModel := getAnthropicModel(model) klog.V(2).Infof("Starting new Anthropic chat session with model: %s", selectedModel) return &anthropicChatSession{ client: c.client, model: selectedModel, systemPrompt: systemPrompt, messages: []anthropic.MessageParam{}, promptCaching: anthropicPromptCaching, extendedThinking: anthropicExtendedThinking, } } // GenerateCompletion generates a single completion for the given request. func (c *AnthropicClient) GenerateCompletion(ctx context.Context, req *CompletionRequest) (CompletionResponse, error) { chat := c.StartChat("", req.Model) chatResponse, err := chat.Send(ctx, req.Prompt) if err != nil { return nil, err } return &anthropicCompletionResponse{chatResponse: chatResponse}, nil } // SetResponseSchema is not supported by the native Anthropic provider. func (c *AnthropicClient) SetResponseSchema(schema *Schema) error { klog.Warning("AnthropicClient.SetResponseSchema is not supported by the native Anthropic provider") return nil } // ListModels returns the list of supported Anthropic Claude models. func (c *AnthropicClient) ListModels(ctx context.Context) ([]string, error) { return []string{ // Claude 4.6 (latest) "claude-opus-4-6", "claude-sonnet-4-6", // Claude 4.5 "claude-opus-4-5", "claude-opus-4-5-20251101", "claude-sonnet-4-5", "claude-sonnet-4-5-20250929", "claude-haiku-4-5", "claude-haiku-4-5-20251001", // Claude 3.7 "claude-3-7-sonnet-latest", "claude-3-7-sonnet-20250219", // Claude 3.5 "claude-3-5-haiku-latest", "claude-3-5-haiku-20241022", // Claude 3 "claude-3-opus-latest", "claude-3-opus-20240229", "claude-3-haiku-20240307", }, nil } // anthropicChatSession implements the Chat interface for Anthropic conversations. type anthropicChatSession struct { client *anthropic.Client model string systemPrompt string messages []anthropic.MessageParam tools []anthropic.ToolUnionParam promptCaching bool extendedThinking bool } // Ensure anthropicChatSession implements the Chat interface. var _ Chat = (*anthropicChatSession)(nil) // Initialize initializes the chat with a previous conversation history. func (c *anthropicChatSession) Initialize(history []*api.Message) error { c.messages = make([]anthropic.MessageParam, 0, len(history)) for _, msg := range history { var role anthropic.MessageParamRole switch msg.Source { case api.MessageSourceUser: role = anthropic.MessageParamRoleUser case api.MessageSourceModel, api.MessageSourceAgent: role = anthropic.MessageParamRoleAssistant default: continue } if msg.Type != api.MessageTypeText || msg.Payload == nil { continue } var content string if textPayload, ok := msg.Payload.(string); ok { content = textPayload } else { content = fmt.Sprintf("%v", msg.Payload) } if content == "" { continue } param := anthropic.MessageParam{ Role: role, Content: []anthropic.ContentBlockParamUnion{anthropic.NewTextBlock(content)}, } c.messages = append(c.messages, param) } return nil } // SetFunctionDefinitions configures the available functions for tool use. func (c *anthropicChatSession) SetFunctionDefinitions(functions []*FunctionDefinition) error { c.tools = nil if len(functions) == 0 { return nil } c.tools = make([]anthropic.ToolUnionParam, len(functions)) for i, fn := range functions { // Build input schema properties from gollm Schema var properties any var required []string if fn.Parameters != nil { schemaBytes, err := json.Marshal(fn.Parameters) if err != nil { return fmt.Errorf("failed to marshal parameters for function %s: %w", fn.Name, err) } var schemaMap map[string]any if err := json.Unmarshal(schemaBytes, &schemaMap); err != nil { return fmt.Errorf("failed to unmarshal parameters for function %s: %w", fn.Name, err) } properties = schemaMap["properties"] if reqSlice, ok := schemaMap["required"].([]any); ok { for _, r := range reqSlice { if s, ok := r.(string); ok { required = append(required, s) } } } } tool := anthropic.ToolParam{ Name: fn.Name, Description: anthropic.String(fn.Description), InputSchema: anthropic.ToolInputSchemaParam{ Properties: properties, Required: required, }, } // Apply prompt cache breakpoint to the last tool definition if c.promptCaching && i == len(functions)-1 { tool.CacheControl = anthropic.NewCacheControlEphemeralParam() } c.tools[i] = anthropic.ToolUnionParam{OfTool: &tool} } return nil } // buildSystemBlocks constructs the system prompt blocks, optionally with cache_control. func (c *anthropicChatSession) buildSystemBlocks() []anthropic.TextBlockParam { if c.systemPrompt == "" { return nil } block := anthropic.TextBlockParam{ Text: c.systemPrompt, } if c.promptCaching { block.CacheControl = anthropic.NewCacheControlEphemeralParam() } return []anthropic.TextBlockParam{block} } // addContentsToHistory processes and appends user messages to chat history. func (c *anthropicChatSession) addContentsToHistory(contents []any) error { var blocks []anthropic.ContentBlockParamUnion for _, content := range contents { switch v := content.(type) { case string: blocks = append(blocks, anthropic.NewTextBlock(v)) case FunctionCallResult: resultJSON, err := json.Marshal(v.Result) if err != nil { return fmt.Errorf("failed to marshal function call result %q: %w", v.Name, err) } // Detect error from result map isError := false if v.Result != nil { if errVal, ok := v.Result["error"]; ok { if errBool, isBool := errVal.(bool); isBool && errBool { isError = true } } if statusVal, ok := v.Result["status"]; ok { if statusStr, isStr := statusVal.(string); isStr && (statusStr == "failed" || statusStr == "error") { isError = true } } } blocks = append(blocks, anthropic.NewToolResultBlock(v.ID, string(resultJSON), isError)) default: return fmt.Errorf("unhandled content type: %T", content) } } if len(blocks) > 0 { c.messages = append(c.messages, anthropic.NewUserMessage(blocks...)) } return nil } // Send sends a message and returns a non-streaming response. func (c *anthropicChatSession) Send(ctx context.Context, contents ...any) (ChatResponse, error) { if len(contents) == 0 { return nil, errors.New("no content provided") } if err := c.addContentsToHistory(contents); err != nil { return nil, err } const thinkingBudget = 8000 maxTokens := anthropicMaxTokens if c.extendedThinking { // max_tokens must exceed budget_tokens maxTokens = thinkingBudget + anthropicMaxTokens } params := anthropic.MessageNewParams{ Model: anthropic.Model(c.model), MaxTokens: maxTokens, Messages: c.messages, System: c.buildSystemBlocks(), } if len(c.tools) > 0 { params.Tools = c.tools } if c.extendedThinking { params.Thinking = anthropic.ThinkingConfigParamOfEnabled(thinkingBudget) } msg, err := c.client.Messages.New(ctx, params) if err != nil { return nil, fmt.Errorf("anthropic message error: %w", err) } // Append assistant message to history c.messages = append(c.messages, msg.ToParam()) return &anthropicResponse{message: msg}, nil } // SendStreaming sends a message and returns a streaming response iterator. func (c *anthropicChatSession) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { if len(contents) == 0 { return nil, errors.New("no content provided") } if err := c.addContentsToHistory(contents); err != nil { return nil, err } const thinkingBudget = 8000 maxTokens := anthropicMaxTokens if c.extendedThinking { // max_tokens must exceed budget_tokens maxTokens = thinkingBudget + anthropicMaxTokens } params := anthropic.MessageNewParams{ Model: anthropic.Model(c.model), MaxTokens: maxTokens, Messages: c.messages, System: c.buildSystemBlocks(), } if len(c.tools) > 0 { params.Tools = c.tools } if c.extendedThinking { params.Thinking = anthropic.ThinkingConfigParamOfEnabled(thinkingBudget) } stream := c.client.Messages.NewStreaming(ctx, params) return func(yield func(ChatResponse, error) bool) { defer stream.Close() // Accumulated message for history acc := anthropic.Message{} // Tool accumulation state type partialTool struct { id string name string input strings.Builder } toolsByIndex := make(map[int64]*partialTool) for stream.Next() { event := stream.Current() // Accumulate for history if err := acc.Accumulate(event); err != nil { klog.V(2).Infof("Anthropic accumulate error: %v", err) } switch ev := event.AsAny().(type) { case anthropic.ContentBlockStartEvent: // Register new tool_use block if ev.ContentBlock.Type == "tool_use" { toolsByIndex[ev.Index] = &partialTool{ id: ev.ContentBlock.ID, name: ev.ContentBlock.Name, } } case anthropic.ContentBlockDeltaEvent: switch delta := ev.Delta.AsAny().(type) { case anthropic.TextDelta: if !yield(&anthropicStreamResponse{text: delta.Text}, nil) { return } case anthropic.InputJSONDelta: // Accumulate tool input JSON if pt, ok := toolsByIndex[ev.Index]; ok { pt.input.WriteString(delta.PartialJSON) } case anthropic.ThinkingDelta: // thinking content is kept in history via accumulator, not yielded to UI } case anthropic.ContentBlockStopEvent: // Check if a tool_use block completed if pt, ok := toolsByIndex[ev.Index]; ok { inputJSON := pt.input.String() var args map[string]any if inputJSON != "" { if err := json.Unmarshal([]byte(inputJSON), &args); err != nil { klog.V(2).Infof("Failed to unmarshal tool input: %v", err) args = make(map[string]any) } } else { args = make(map[string]any) } fc := FunctionCall{ ID: pt.id, Name: pt.name, Arguments: args, } if !yield(&anthropicStreamResponse{functionCall: &fc}, nil) { return } delete(toolsByIndex, ev.Index) } } } if err := stream.Err(); err != nil { yield(nil, fmt.Errorf("anthropic stream error: %w", err)) return } // Append accumulated assistant message to history if len(acc.Content) > 0 { c.messages = append(c.messages, acc.ToParam()) } // Yield final usage so callers can observe token/cache counts if acc.Usage.InputTokens > 0 || acc.Usage.OutputTokens > 0 { yield(&anthropicStreamResponse{usage: &acc.Usage}, nil) } }, nil } // IsRetryableError determines if an error from the Anthropic API should be retried. func (c *anthropicChatSession) IsRetryableError(err error) bool { if err == nil { return false } var apiErr *anthropic.Error if errors.As(err, &apiErr) { sc := apiErr.StatusCode // 429 = rate limit, 529 = overloaded, 5xx = server errors return sc == 429 || sc == 529 || (sc >= 500 && sc < 600) } return DefaultIsRetryableError(err) } // anthropicResponse implements ChatResponse for non-streaming responses. type anthropicResponse struct { message *anthropic.Message } var _ ChatResponse = (*anthropicResponse)(nil) func (r *anthropicResponse) UsageMetadata() any { if r.message != nil { return r.message.Usage } return nil } func (r *anthropicResponse) Candidates() []Candidate { if r.message == nil { return nil } return []Candidate{&anthropicCandidate{content: r.message.Content}} } // anthropicStreamResponse implements ChatResponse for streaming responses. type anthropicStreamResponse struct { text string functionCall *FunctionCall usage *anthropic.Usage } var _ ChatResponse = (*anthropicStreamResponse)(nil) func (r *anthropicStreamResponse) UsageMetadata() any { if r.usage != nil { return r.usage } return nil } func (r *anthropicStreamResponse) Candidates() []Candidate { if r.text == "" && r.functionCall == nil { return nil } return []Candidate{&anthropicStreamCandidate{ text: r.text, functionCall: r.functionCall, }} } // anthropicCandidate implements Candidate for non-streaming responses. type anthropicCandidate struct { content []anthropic.ContentBlockUnion } var _ Candidate = (*anthropicCandidate)(nil) func (c *anthropicCandidate) String() string { var sb strings.Builder for _, block := range c.content { switch block.Type { case "text": tb := block.AsText() sb.WriteString(tb.Text) } } return sb.String() } func (c *anthropicCandidate) Parts() []Part { var parts []Part for _, block := range c.content { switch block.Type { case "text": tb := block.AsText() if tb.Text != "" { parts = append(parts, &anthropicTextPart{text: tb.Text}) } case "tool_use": tu := block.AsToolUse() var args map[string]any if len(tu.Input) > 0 { if err := json.Unmarshal(tu.Input, &args); err != nil { klog.V(2).Infof("Failed to unmarshal tool input: %v", err) args = make(map[string]any) } } else { args = make(map[string]any) } parts = append(parts, &anthropicToolPart{ functionCall: FunctionCall{ ID: tu.ID, Name: tu.Name, Arguments: args, }, }) case "thinking": // ThinkingBlock — do not yield to UI, skip } } return parts } // anthropicStreamCandidate implements Candidate for streaming responses. type anthropicStreamCandidate struct { text string functionCall *FunctionCall } var _ Candidate = (*anthropicStreamCandidate)(nil) func (c *anthropicStreamCandidate) String() string { if c.text != "" { return c.text } if c.functionCall != nil { return fmt.Sprintf("FunctionCall(%s)", c.functionCall.Name) } return "" } func (c *anthropicStreamCandidate) Parts() []Part { var parts []Part if c.text != "" { parts = append(parts, &anthropicTextPart{text: c.text}) } if c.functionCall != nil { parts = append(parts, &anthropicToolPart{functionCall: *c.functionCall}) } return parts } // anthropicTextPart implements Part for text content. type anthropicTextPart struct { text string } var _ Part = (*anthropicTextPart)(nil) func (p *anthropicTextPart) AsText() (string, bool) { return p.text, p.text != "" } func (p *anthropicTextPart) AsFunctionCalls() ([]FunctionCall, bool) { return nil, false } // anthropicToolPart implements Part for tool/function calls. type anthropicToolPart struct { functionCall FunctionCall } var _ Part = (*anthropicToolPart)(nil) func (p *anthropicToolPart) AsText() (string, bool) { return "", false } func (p *anthropicToolPart) AsFunctionCalls() ([]FunctionCall, bool) { return []FunctionCall{p.functionCall}, true } // anthropicCompletionResponse wraps a ChatResponse to implement CompletionResponse. type anthropicCompletionResponse struct { chatResponse ChatResponse } var _ CompletionResponse = (*anthropicCompletionResponse)(nil) func (r *anthropicCompletionResponse) Response() string { if r.chatResponse == nil { return "" } candidates := r.chatResponse.Candidates() if len(candidates) == 0 { return "" } for _, part := range candidates[0].Parts() { if text, ok := part.AsText(); ok { return text } } return "" } func (r *anthropicCompletionResponse) UsageMetadata() any { if r.chatResponse == nil { return nil } return r.chatResponse.UsageMetadata() } func getAnthropicModel(model string) string { if model != "" && strings.HasPrefix(model, "claude") { klog.V(4).Infof("Using explicitly provided Anthropic model: %s", model) return model } if model != "" { klog.V(2).Infof("Ignoring non-Claude model %q, falling back to default", model) } if anthropicDefaultModel != "" { klog.V(2).Infof("Using Anthropic model from environment: %s", anthropicDefaultModel) return anthropicDefaultModel } defaultModel := "claude-sonnet-4-6" klog.V(2).Infof("Using default Anthropic model: %s", defaultModel) return defaultModel } ================================================ FILE: gollm/anthropic_test.go ================================================ // Copyright 2026 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "fmt" "net/http" "strconv" "testing" anthropic "github.com/anthropics/anthropic-sdk-go" ) // TestAnthropicProviderRegistration verifies that the "anthropic" provider // is registered in the global registry when the package initializes. func TestAnthropicProviderRegistration(t *testing.T) { providers := globalRegistry.listProviders() found := false for _, p := range providers { if p == "anthropic" { found = true break } } if !found { t.Errorf("expected 'anthropic' to be registered, got providers: %v", providers) } } // TestAnthropicAddContentsToHistory verifies that string and FunctionCallResult // contents are correctly converted to Anthropic MessageParam history entries. func TestAnthropicAddContentsToHistory(t *testing.T) { tests := []struct { name string contents []any wantMsgs int wantRole anthropic.MessageParamRole wantErr bool }{ { name: "string content creates user message", contents: []any{"hello world"}, wantMsgs: 1, wantRole: anthropic.MessageParamRoleUser, wantErr: false, }, { name: "FunctionCallResult creates user message with tool_result block", contents: []any{FunctionCallResult{ ID: "tool_123", Name: "kubectl", Result: map[string]any{"output": "pods running"}, }}, wantMsgs: 1, wantRole: anthropic.MessageParamRoleUser, wantErr: false, }, { name: "unhandled content type returns error", contents: []any{12345}, wantMsgs: 0, wantErr: true, }, { name: "empty contents adds no messages", contents: []any{}, wantMsgs: 0, wantErr: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { session := &anthropicChatSession{ messages: []anthropic.MessageParam{}, promptCaching: false, } err := session.addContentsToHistory(tt.contents) if tt.wantErr { if err == nil { t.Error("expected error but got none") } return } if err != nil { t.Errorf("unexpected error: %v", err) return } if len(session.messages) != tt.wantMsgs { t.Errorf("expected %d messages, got %d", tt.wantMsgs, len(session.messages)) return } if tt.wantMsgs > 0 { if session.messages[0].Role != tt.wantRole { t.Errorf("expected role %q, got %q", tt.wantRole, session.messages[0].Role) } } }) } } // TestAnthropicBuildSystemBlocks_WithCaching verifies that the system prompt // block includes cache_control when prompt caching is enabled. func TestAnthropicBuildSystemBlocks_WithCaching(t *testing.T) { session := &anthropicChatSession{ systemPrompt: "You are a helpful Kubernetes assistant.", promptCaching: true, } blocks := session.buildSystemBlocks() if len(blocks) != 1 { t.Fatalf("expected 1 system block, got %d", len(blocks)) } block := blocks[0] if block.Text != session.systemPrompt { t.Errorf("expected text %q, got %q", session.systemPrompt, block.Text) } // With caching enabled, CacheControl should be set (non-zero type field) if block.CacheControl.Type == "" { t.Error("expected CacheControl.Type to be set when prompt caching is enabled") } } // TestAnthropicBuildSystemBlocks_WithoutCaching verifies that the system prompt // block does NOT include cache_control when prompt caching is disabled. func TestAnthropicBuildSystemBlocks_WithoutCaching(t *testing.T) { session := &anthropicChatSession{ systemPrompt: "You are a helpful Kubernetes assistant.", promptCaching: false, } blocks := session.buildSystemBlocks() if len(blocks) != 1 { t.Fatalf("expected 1 system block, got %d", len(blocks)) } block := blocks[0] if block.Text != session.systemPrompt { t.Errorf("expected text %q, got %q", session.systemPrompt, block.Text) } // Without caching, CacheControl.Type should be empty/zero if block.CacheControl.Type != "" { t.Error("expected CacheControl to be empty when prompt caching is disabled") } } // TestAnthropicBuildSystemBlocks_EmptyPrompt verifies that an empty system // prompt returns nil blocks. func TestAnthropicBuildSystemBlocks_EmptyPrompt(t *testing.T) { session := &anthropicChatSession{ systemPrompt: "", promptCaching: true, } blocks := session.buildSystemBlocks() if blocks != nil { t.Errorf("expected nil blocks for empty system prompt, got %v", blocks) } } // TestAnthropicSetFunctionDefinitions_CacheControl verifies that when prompt // caching is enabled, the cache breakpoint is placed on the last tool definition. func TestAnthropicSetFunctionDefinitions_CacheControl(t *testing.T) { session := &anthropicChatSession{ promptCaching: true, } functions := []*FunctionDefinition{ { Name: "kubectl", Description: "Run a kubectl command", Parameters: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "command": {Type: TypeString, Description: "The kubectl command"}, }, }, }, { Name: "bash", Description: "Run a bash command", Parameters: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "command": {Type: TypeString, Description: "The bash command"}, }, }, }, } if err := session.SetFunctionDefinitions(functions); err != nil { t.Fatalf("SetFunctionDefinitions error: %v", err) } if len(session.tools) != 2 { t.Fatalf("expected 2 tools, got %d", len(session.tools)) } // The first tool should NOT have cache_control firstTool := session.tools[0].OfTool if firstTool == nil { t.Fatal("expected first tool to be non-nil") } if firstTool.CacheControl.Type != "" { t.Error("expected first tool to NOT have cache_control") } // The last tool SHOULD have cache_control lastTool := session.tools[1].OfTool if lastTool == nil { t.Fatal("expected last tool to be non-nil") } if lastTool.CacheControl.Type == "" { t.Error("expected last tool to have cache_control set") } } // TestAnthropicSetFunctionDefinitions_NoCacheControl verifies that when prompt // caching is disabled, no tool gets a cache breakpoint. func TestAnthropicSetFunctionDefinitions_NoCacheControl(t *testing.T) { session := &anthropicChatSession{ promptCaching: false, } functions := []*FunctionDefinition{ {Name: "kubectl", Description: "Run kubectl"}, {Name: "bash", Description: "Run bash"}, } if err := session.SetFunctionDefinitions(functions); err != nil { t.Fatalf("SetFunctionDefinitions error: %v", err) } for i, tool := range session.tools { if tool.OfTool == nil { t.Fatalf("tool[%d] is nil", i) } if tool.OfTool.CacheControl.Type != "" { t.Errorf("tool[%d] should NOT have cache_control when caching is disabled", i) } } } // TestAnthropicSetFunctionDefinitions_EmptyList verifies that setting an empty // function list clears any previously set tools. func TestAnthropicSetFunctionDefinitions_EmptyList(t *testing.T) { session := &anthropicChatSession{ tools: []anthropic.ToolUnionParam{{OfTool: &anthropic.ToolParam{Name: "old"}}}, } if err := session.SetFunctionDefinitions([]*FunctionDefinition{}); err != nil { t.Fatalf("unexpected error: %v", err) } if len(session.tools) != 0 { t.Errorf("expected tools to be cleared, got %d tools", len(session.tools)) } } // TestAnthropicIsRetryableError verifies the retry classification logic. func TestAnthropicIsRetryableError(t *testing.T) { session := &anthropicChatSession{} tests := []struct { name string err error wantRetry bool }{ { name: "nil error is not retryable", err: nil, wantRetry: false, }, { name: "rate limit 429 is retryable", err: makeAnthropicAPIError(429), wantRetry: true, }, { name: "overloaded 529 is retryable", err: makeAnthropicAPIError(529), wantRetry: true, }, { name: "internal server error 500 is retryable", err: makeAnthropicAPIError(500), wantRetry: true, }, { name: "bad gateway 502 is retryable", err: makeAnthropicAPIError(502), wantRetry: true, }, { name: "bad request 400 is not retryable", err: makeAnthropicAPIError(400), wantRetry: false, }, { name: "not found 404 is not retryable", err: makeAnthropicAPIError(404), wantRetry: false, }, { name: "authentication error 401 is not retryable", err: makeAnthropicAPIError(401), wantRetry: false, }, { name: "plain error falls through to default", err: fmt.Errorf("some generic error"), wantRetry: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := session.IsRetryableError(tt.err) if got != tt.wantRetry { t.Errorf("IsRetryableError(%v) = %v, want %v", tt.err, got, tt.wantRetry) } }) } } // makeAnthropicAPIError creates a fake *anthropic.Error with the given status code // for testing IsRetryableError. anthropic.Error is a public alias for the internal // apierror.Error struct, so we construct it via the alias directly. func makeAnthropicAPIError(statusCode int) error { return &anthropic.Error{ StatusCode: statusCode, Request: &http.Request{Method: "POST"}, Response: &http.Response{StatusCode: statusCode}, } } // TestAnthropicStreamResponseUsageMetadata verifies that UsageMetadata returns // non-nil when a usage struct is set, and nil otherwise. func TestAnthropicStreamResponseUsageMetadata(t *testing.T) { t.Run("returns nil when no usage set", func(t *testing.T) { r := &anthropicStreamResponse{text: "hello"} if r.UsageMetadata() != nil { t.Error("expected nil UsageMetadata for text-only response") } }) t.Run("returns usage when set", func(t *testing.T) { usage := &anthropic.Usage{InputTokens: 10, OutputTokens: 20} r := &anthropicStreamResponse{usage: usage} got := r.UsageMetadata() if got == nil { t.Fatal("expected non-nil UsageMetadata") } u, ok := got.(*anthropic.Usage) if !ok { t.Fatalf("expected *anthropic.Usage, got %T", got) } if u.InputTokens != 10 || u.OutputTokens != 20 { t.Errorf("unexpected usage values: %+v", u) } }) t.Run("usage-only response returns nil candidates", func(t *testing.T) { usage := &anthropic.Usage{InputTokens: 5, OutputTokens: 15} r := &anthropicStreamResponse{usage: usage} if r.Candidates() != nil { t.Error("expected nil Candidates for usage-only response") } }) } // TestAnthropicMaxTokensDefault verifies that the package-level default is 4096. func TestAnthropicMaxTokensDefault(t *testing.T) { // Save and restore orig := anthropicMaxTokens defer func() { anthropicMaxTokens = orig }() anthropicMaxTokens = 4096 if anthropicMaxTokens != 4096 { t.Errorf("expected default max tokens 4096, got %d", anthropicMaxTokens) } } // TestAnthropicMaxTokensEnvVar verifies that ANTHROPIC_MAX_TOKENS is parsed // and applied to the package-level variable. func TestAnthropicMaxTokensEnvVar(t *testing.T) { orig := anthropicMaxTokens defer func() { anthropicMaxTokens = orig }() t.Run("valid value is applied", func(t *testing.T) { t.Setenv("ANTHROPIC_MAX_TOKENS", "2048") // Simulate what init() does anthropicMaxTokens = 4096 if v := t.TempDir(); v != "" { // just to use t } // Re-run the parsing logic inline (mirrors init()) if v := "2048"; v != "" { if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { anthropicMaxTokens = n } } if anthropicMaxTokens != 2048 { t.Errorf("expected 2048, got %d", anthropicMaxTokens) } }) t.Run("zero value is rejected, default kept", func(t *testing.T) { anthropicMaxTokens = 4096 if v := "0"; v != "" { if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { anthropicMaxTokens = n } } if anthropicMaxTokens != 4096 { t.Errorf("expected default 4096 to be kept, got %d", anthropicMaxTokens) } }) t.Run("negative value is rejected, default kept", func(t *testing.T) { anthropicMaxTokens = 4096 if v := "-100"; v != "" { if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { anthropicMaxTokens = n } } if anthropicMaxTokens != 4096 { t.Errorf("expected default 4096 to be kept, got %d", anthropicMaxTokens) } }) t.Run("non-numeric value is rejected, default kept", func(t *testing.T) { anthropicMaxTokens = 4096 if v := "abc"; v != "" { if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { anthropicMaxTokens = n } } if anthropicMaxTokens != 4096 { t.Errorf("expected default 4096 to be kept, got %d", anthropicMaxTokens) } }) } // TestGetAnthropicModel verifies the model selection priority. func TestGetAnthropicModel(t *testing.T) { // Save and restore original values origDefault := anthropicDefaultModel defer func() { anthropicDefaultModel = origDefault }() t.Run("explicit model is highest priority", func(t *testing.T) { anthropicDefaultModel = "env-model" got := getAnthropicModel("explicit-model") if got != "explicit-model" { t.Errorf("expected 'explicit-model', got %q", got) } }) t.Run("env var used when no explicit model", func(t *testing.T) { anthropicDefaultModel = "env-model" got := getAnthropicModel("") if got != "env-model" { t.Errorf("expected 'env-model', got %q", got) } }) t.Run("default used when no explicit model and no env var", func(t *testing.T) { anthropicDefaultModel = "" got := getAnthropicModel("") if got != "claude-sonnet-4-6" { t.Errorf("expected 'claude-sonnet-4-6', got %q", got) } }) } ================================================ FILE: gollm/azopenai.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "fmt" "os" "slices" "strings" "k8s.io/klog/v2" "github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai" "github.com/Azure/azure-sdk-for-go/sdk/azcore" "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" "github.com/Azure/azure-sdk-for-go/sdk/azidentity" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) func init() { if err := RegisterProvider("azopenai", azureOpenAIFactory); err != nil { klog.Fatalf("Failed to register azopenai provider: %v", err) } } /* azureOpenAIFactory is the provider factory function for Azure OpenAI. Supports ClientOptions for custom configuration. */ func azureOpenAIFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewAzureOpenAIClient(ctx, opts) } type AzureOpenAIClient struct { client *azopenai.Client endpoint string } var _ Client = &AzureOpenAIClient{} // NewAzureOpenAIClient creates a new Azure OpenAI client. // Supports ClientOptions and SkipVerifySSL for custom HTTP transport. func NewAzureOpenAIClient(ctx context.Context, opts ClientOptions) (*AzureOpenAIClient, error) { azureOpenAIEndpoint := os.Getenv("AZURE_OPENAI_ENDPOINT") if opts.URL != nil && opts.URL.Host != "" { opts.URL.Scheme = "https" azureOpenAIEndpoint = opts.URL.String() } if azureOpenAIEndpoint == "" { return nil, fmt.Errorf("AZURE_OPENAI_ENDPOINT environment variable not set") } azureOpenAIClient := AzureOpenAIClient{ endpoint: azureOpenAIEndpoint, } // Create a custom HTTP client (supports SkipVerifySSL) httpClient := createCustomHTTPClient(opts.SkipVerifySSL) azureOpenAIKey := os.Getenv("AZURE_OPENAI_API_KEY") clientOpts := &azopenai.ClientOptions{ ClientOptions: azcore.ClientOptions{ Transport: httpClient, }, } if azureOpenAIKey != "" { keyCredential := azcore.NewKeyCredential(azureOpenAIKey) client, err := azopenai.NewClientWithKeyCredential(azureOpenAIEndpoint, keyCredential, clientOpts) if err != nil { return nil, fmt.Errorf("failed to create azure openai client: %w", err) } azureOpenAIClient.client = client } else { credential, err := azidentity.NewDefaultAzureCredential(nil) if err != nil { return nil, fmt.Errorf("failed to get credential: %w", err) } client, err := azopenai.NewClient(azureOpenAIEndpoint, credential, clientOpts) if err != nil { return nil, fmt.Errorf("failed to create azure openai client: %w", err) } azureOpenAIClient.client = client } return &azureOpenAIClient, nil } func (c *AzureOpenAIClient) Close() error { return nil } func (c *AzureOpenAIClient) GenerateCompletion(ctx context.Context, request *CompletionRequest) (CompletionResponse, error) { req := azopenai.ChatCompletionsOptions{ Messages: []azopenai.ChatRequestMessageClassification{ &azopenai.ChatRequestUserMessage{Content: azopenai.NewChatRequestUserMessageContent(request.Prompt)}, }, DeploymentName: &request.Model, } resp, err := c.client.GetChatCompletions(ctx, req, nil) if err != nil { return nil, err } if len(resp.Choices) == 0 || resp.Choices[0].Message == nil || resp.Choices[0].Message.Content == nil { return nil, fmt.Errorf("invalid completion response: %v", resp) } return &AzureOpenAICompletionResponse{response: *resp.Choices[0].Message.Content}, nil } func (c *AzureOpenAIClient) ListModels(ctx context.Context) ([]string, error) { cred, err := azidentity.NewDefaultAzureCredential(nil) if err != nil { return nil, fmt.Errorf("failed to get credential: %w", err) } subClient, err := armsubscription.NewSubscriptionsClient(cred, nil) if err != nil { return nil, fmt.Errorf("failed to create subscriptions client: %w", err) } subPager := subClient.NewListPager(nil) for subPager.More() { subResp, err := subPager.NextPage(context.Background()) if err != nil { return nil, fmt.Errorf("failed to get subscriptions page: %w", err) } for _, sub := range subResp.Value { accountClient, err := armcognitiveservices.NewAccountsClient(*sub.SubscriptionID, cred, nil) if err != nil { return nil, fmt.Errorf("failed to create accounts client: %w", err) } accountPager := accountClient.NewListPager(nil) for accountPager.More() { accountResp, err := accountPager.NextPage(context.Background()) if err != nil { return nil, fmt.Errorf("failed to to get accounts page: %w", err) } for _, account := range accountResp.Value { if account.Kind == nil || !slices.Contains([]string{"OpenAI", "CognitiveServices", "AIServices"}, *account.Kind) { // Not an Azure OpenAI service continue } if account.Properties == nil || account.Properties.Endpoint == nil || strings.TrimSuffix(*account.Properties.Endpoint, "/") != c.endpoint { // Not the expected endpoint continue } resourceID, err := arm.ParseResourceID(*account.ID) if err != nil { return nil, fmt.Errorf("failed to parse resource ID %q: %w", *account.Name, err) } deploymentClient, err := armcognitiveservices.NewDeploymentsClient(*sub.SubscriptionID, cred, nil) if err != nil { return nil, fmt.Errorf("failed to create deployments client: %w", err) } var modelNames []string deploymentPager := deploymentClient.NewListPager(resourceID.ResourceGroupName, *account.Name, nil) for deploymentPager.More() { deploymentResp, err := deploymentPager.NextPage(context.Background()) if err != nil { return nil, fmt.Errorf("failed to get deployments page: %w", err) } for _, deployment := range deploymentResp.Value { modelNames = append(modelNames, *deployment.Name) } } slices.Sort(modelNames) return modelNames, nil } } } } return nil, nil } func (c *AzureOpenAIClient) SetResponseSchema(schema *Schema) error { return nil } func (c *AzureOpenAIClient) StartChat(systemPrompt string, model string) Chat { return &AzureOpenAIChat{ client: c.client, model: model, history: []azopenai.ChatRequestMessageClassification{ &azopenai.ChatRequestSystemMessage{Content: azopenai.NewChatRequestSystemMessageContent(systemPrompt)}, }, } } type AzureOpenAICompletionResponse struct { response string } func (r *AzureOpenAICompletionResponse) Response() string { return r.response } func (r *AzureOpenAICompletionResponse) UsageMetadata() any { return nil } type AzureOpenAIChat struct { client *azopenai.Client model string history []azopenai.ChatRequestMessageClassification tools []azopenai.ChatCompletionsToolDefinitionClassification } func (c *AzureOpenAIChat) Send(ctx context.Context, contents ...any) (ChatResponse, error) { for _, content := range contents { switch v := content.(type) { case string: message := azopenai.ChatRequestUserMessage{ Content: azopenai.NewChatRequestUserMessageContent(v), } c.history = append(c.history, &message) case FunctionCallResult: message := azopenai.ChatRequestUserMessage{ Content: azopenai.NewChatRequestUserMessageContent(fmt.Sprintf("Function call result: %s", v.Result)), } c.history = append(c.history, &message) default: return nil, fmt.Errorf("unsupported content type: %T", v) } } resp, err := c.client.GetChatCompletions(ctx, azopenai.ChatCompletionsOptions{ DeploymentName: &c.model, Messages: c.history, Tools: c.tools, }, nil) if err != nil { return nil, err } if len(resp.Choices) == 0 { return nil, fmt.Errorf("no response from Azure OpenAI: %v", resp) } return &AzureOpenAIChatResponse{azureOpenAIResponse: resp}, nil } func (c *AzureOpenAIChat) IsRetryableError(err error) bool { // TODO: Implement this return false } func (c *AzureOpenAIChat) Initialize(messages []*api.Message) error { klog.Warning("chat history persistence is not supported for provider 'azopenai', using in-memory chat history") return nil } func (c *AzureOpenAIChat) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { // TODO: Implement streaming response, err := c.Send(ctx, contents...) if err != nil { return nil, err } return singletonChatResponseIterator(response), nil } type AzureOpenAIChatResponse struct { azureOpenAIResponse azopenai.GetChatCompletionsResponse } var _ ChatResponse = &AzureOpenAIChatResponse{} func (r *AzureOpenAIChatResponse) MarshalJSON() ([]byte, error) { formatted := RecordChatResponse{ Raw: r.azureOpenAIResponse, } return json.Marshal(&formatted) } func (r *AzureOpenAIChatResponse) String() string { return fmt.Sprintf("AzureOpenAIChatResponse{candidates=%v}", r.azureOpenAIResponse.Choices) } func (r *AzureOpenAIChatResponse) UsageMetadata() any { return r.azureOpenAIResponse.Usage } func (r *AzureOpenAIChatResponse) Candidates() []Candidate { var candidates []Candidate for _, candidate := range r.azureOpenAIResponse.Choices { candidates = append(candidates, &AzureOpenAICandidate{candidate: candidate}) } return candidates } type AzureOpenAICandidate struct { candidate azopenai.ChatChoice } func (r *AzureOpenAICandidate) String() string { var response strings.Builder response.WriteString("[") for i, parts := range r.Parts() { if i > 0 { response.WriteString(", ") } text, ok := parts.AsText() if ok { response.WriteString(text) } functionCalls, ok := parts.AsFunctionCalls() if ok { response.WriteString("functionCalls=[") for _, functionCall := range functionCalls { response.WriteString(fmt.Sprintf("%q(args=%v)", functionCall.Name, functionCall.Arguments)) } response.WriteString("]}") } } response.WriteString("]}") return response.String() } func (r *AzureOpenAICandidate) Parts() []Part { var parts []Part if r.candidate.Message != nil { parts = append(parts, &AzureOpenAIPart{ text: r.candidate.Message.Content, }) } for _, tool := range r.candidate.Message.ToolCalls { if tool == nil { continue } parts = append(parts, &AzureOpenAIPart{ functionCall: tool.(*azopenai.ChatCompletionsFunctionToolCall).Function, }) } return parts } type AzureOpenAIPart struct { text *string functionCall *azopenai.FunctionCall } func (p *AzureOpenAIPart) AsText() (string, bool) { if p.text != nil && len(*p.text) > 0 { return *p.text, true } return "", false } func (p *AzureOpenAIPart) AsFunctionCalls() ([]FunctionCall, bool) { if p.functionCall != nil { argumentsObj := map[string]any{} err := json.Unmarshal([]byte(*p.functionCall.Arguments), &argumentsObj) if err != nil { return nil, false } functionCalls := []FunctionCall{ { Name: *p.functionCall.Name, Arguments: argumentsObj, }, } return functionCalls, true } return nil, false } func (c *AzureOpenAIChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error { var tools []azopenai.ChatCompletionsToolDefinitionClassification for _, functionDefinition := range functionDefinitions { tools = append(tools, &azopenai.ChatCompletionsFunctionToolDefinition{Function: fnDefToAzureOpenAITool(functionDefinition)}) } c.tools = tools return nil } func fnDefToAzureOpenAITool(fnDef *FunctionDefinition) *azopenai.ChatCompletionsFunctionToolDefinitionFunction { properties := make(map[string]any) for paramName, param := range fnDef.Parameters.Properties { properties[paramName] = map[string]any{ "type": string(param.Type), "description": param.Description, } } parameters := map[string]any{ "type": "object", "properties": properties, } if len(fnDef.Parameters.Required) > 0 { parameters["required"] = fnDef.Parameters.Required } jsonBytes, _ := json.Marshal(parameters) tool := azopenai.ChatCompletionsFunctionToolDefinitionFunction{ Name: &fnDef.Name, Description: &fnDef.Description, Parameters: jsonBytes, } return &tool } ================================================ FILE: gollm/bedrock.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "errors" "fmt" "os" "strings" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/document" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" "k8s.io/klog/v2" ) // Register the Bedrock provider factory on package initialization func init() { if err := RegisterProvider("bedrock", newBedrockClientFactory); err != nil { klog.Fatalf("Failed to register bedrock provider: %v", err) } } // newBedrockClientFactory creates a new Bedrock client with the given options func newBedrockClientFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewBedrockClient(ctx, opts) } // BedrockClient implements the gollm.Client interface for AWS Bedrock models type BedrockClient struct { client *bedrockruntime.Client } // Ensure BedrockClient implements the Client interface var _ Client = &BedrockClient{} // NewBedrockClient creates a new client for interacting with AWS Bedrock models func NewBedrockClient(ctx context.Context, opts ClientOptions) (*BedrockClient, error) { // Load AWS config with timeout protection configCtx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() cfg, err := config.LoadDefaultConfig(configCtx) if err != nil { return nil, fmt.Errorf("failed to load AWS config: %w", err) } // Default to us-east-1 for Bedrock if no region is set if cfg.Region == "" { cfg.Region = "us-east-1" } return &BedrockClient{ client: bedrockruntime.NewFromConfig(cfg), }, nil } // Close cleans up any resources used by the client func (c *BedrockClient) Close() error { return nil } // StartChat starts a new chat session with the specified system prompt and model func (c *BedrockClient) StartChat(systemPrompt, model string) Chat { selectedModel := getBedrockModel(model) // Enhance system prompt for tool-use shim compatibility // Detect if tool-use shim is enabled by looking for JSON formatting instructions enhancedPrompt := systemPrompt if strings.Contains(systemPrompt, "```json") && strings.Contains(systemPrompt, "\"action\"") { // Tool-use shim is enabled - add stronger JSON formatting instructions for all Bedrock models enhancedPrompt += "\n\nCRITICAL JSON FORMATTING REQUIREMENTS:\n" enhancedPrompt += "1. You MUST ALWAYS wrap your JSON responses in ```json code blocks exactly as shown in the examples above.\n" enhancedPrompt += "2. NEVER respond with raw JSON without the markdown ```json formatting.\n" enhancedPrompt += "3. Ensure your JSON is syntactically correct with proper commas between fields.\n" enhancedPrompt += "4. This is critical for proper parsing. Example format:\n" enhancedPrompt += "```json\n{\"thought\": \"your reasoning\", \"action\": {\"name\": \"tool_name\", \"command\": \"command\"}}\n```\n" enhancedPrompt += "Note the comma after the \"thought\" field! Malformed JSON will cause failures." } return &bedrockChat{ client: c, systemPrompt: enhancedPrompt, model: selectedModel, messages: []types.Message{}, } } // GenerateCompletion generates a single completion for the given request func (c *BedrockClient) GenerateCompletion(ctx context.Context, req *CompletionRequest) (CompletionResponse, error) { chat := c.StartChat("", req.Model) chatResponse, err := chat.Send(ctx, req.Prompt) if err != nil { return nil, err } // Wrap ChatResponse in a CompletionResponse return &bedrockCompletionResponse{ chatResponse: chatResponse, }, nil } // SetResponseSchema sets the response schema for the client (not supported by Bedrock) func (c *BedrockClient) SetResponseSchema(schema *Schema) error { return fmt.Errorf("response schema not supported by Bedrock") } // ListModels returns the list of supported Bedrock models func (c *BedrockClient) ListModels(ctx context.Context) ([]string, error) { return []string{ "us.anthropic.claude-sonnet-4-20250514-v1:0", // Claude Sonnet 4 (default) "us.anthropic.claude-3-7-sonnet-20250219-v1:0", // Claude 3.7 Sonnet }, nil } // bedrockChat implements the Chat interface for Bedrock conversations type bedrockChat struct { client *BedrockClient systemPrompt string model string messages []types.Message toolConfig *types.ToolConfiguration functionDefs []*FunctionDefinition } func (cs *bedrockChat) Initialize(history []*api.Message) error { cs.messages = make([]types.Message, 0, len(history)) for _, msg := range history { // Convert api.Message to types.Message var role types.ConversationRole switch msg.Source { case api.MessageSourceUser: role = types.ConversationRoleUser case api.MessageSourceModel, api.MessageSourceAgent: role = types.ConversationRoleAssistant default: // Skip unknown message sources continue } // Convert payload to string content var content string if msg.Type == api.MessageTypeText && msg.Payload != nil { if textPayload, ok := msg.Payload.(string); ok { content = textPayload } else { // Try to convert other types to string content = fmt.Sprintf("%v", msg.Payload) } } else { // Skip non-text messages for now continue } if content == "" { continue } bedrockMsg := types.Message{ Role: role, Content: []types.ContentBlock{ &types.ContentBlockMemberText{Value: content}, }, } cs.messages = append(cs.messages, bedrockMsg) } return nil } // Send sends a message to the chat and returns the response func (c *bedrockChat) Send(ctx context.Context, contents ...any) (ChatResponse, error) { if len(contents) == 0 { return nil, errors.New("no content provided") } // Process and append contents to conversation history if err := c.addContentsToHistory(contents); err != nil { return nil, err } // Prepare the request input := &bedrockruntime.ConverseInput{ ModelId: aws.String(c.model), Messages: c.messages, InferenceConfig: &types.InferenceConfiguration{ MaxTokens: aws.Int32(4096), }, } // Add system prompt if provided if c.systemPrompt != "" { input.System = []types.SystemContentBlock{ &types.SystemContentBlockMemberText{Value: c.systemPrompt}, } } // Add tool configuration if functions are defined if c.toolConfig != nil { input.ToolConfig = c.toolConfig } // Call the Bedrock Converse API output, err := c.client.client.Converse(ctx, input) if err != nil { return nil, fmt.Errorf("bedrock converse error: %w", err) } // Extract response content and update conversation history response := &bedrockResponse{ output: output, model: c.model, } // Update conversation history with assistant's response if output.Output != nil { if msg, ok := output.Output.(*types.ConverseOutputMemberMessage); ok { c.messages = append(c.messages, msg.Value) } } return response, nil } // SendStreaming sends a message and returns a streaming response func (c *bedrockChat) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { if len(contents) == 0 { return nil, errors.New("no content provided") } // Process and append contents to conversation history if err := c.addContentsToHistory(contents); err != nil { return nil, err } // Prepare the streaming request input := &bedrockruntime.ConverseStreamInput{ ModelId: aws.String(c.model), Messages: c.messages, InferenceConfig: &types.InferenceConfiguration{ MaxTokens: aws.Int32(4096), }, } // Add system prompt if provided if c.systemPrompt != "" { input.System = []types.SystemContentBlock{ &types.SystemContentBlockMemberText{Value: c.systemPrompt}, } } // Add tool configuration if functions are defined if c.toolConfig != nil { input.ToolConfig = c.toolConfig } // Start the streaming request output, err := c.client.client.ConverseStream(ctx, input) if err != nil { return nil, fmt.Errorf("bedrock stream error: %w", err) } // Return streaming iterator return func(yield func(ChatResponse, error) bool) { defer func() { if stream := output.GetStream(); stream != nil { stream.Close() } }() var assistantMessage types.Message assistantMessage.Role = types.ConversationRoleAssistant var fullContent strings.Builder // Tool state tracking for streaming type partialTool struct { id string name string input strings.Builder } partialTools := make(map[int32]*partialTool) var completedTools []types.ToolUseBlock // Process streaming events stream := output.GetStream() for event := range stream.Events() { switch v := event.(type) { case *types.ConverseStreamOutputMemberContentBlockDelta: // Handle text deltas if textDelta, ok := v.Value.Delta.(*types.ContentBlockDeltaMemberText); ok { fullContent.WriteString(textDelta.Value) response := &bedrockStreamResponse{ content: textDelta.Value, model: c.model, done: false, } if !yield(response, nil) { return } } // Handle tool input deltas if toolDelta, ok := v.Value.Delta.(*types.ContentBlockDeltaMemberToolUse); ok { idx := aws.ToInt32(v.Value.ContentBlockIndex) if partial, exists := partialTools[idx]; exists { deltaInput := aws.ToString(toolDelta.Value.Input) partial.input.WriteString(deltaInput) } } case *types.ConverseStreamOutputMemberContentBlockStart: // Handle content block start (for tool calls) if v.Value.Start != nil { if toolStart, ok := v.Value.Start.(*types.ContentBlockStartMemberToolUse); ok { // Store partial tool for input accumulation idx := aws.ToInt32(v.Value.ContentBlockIndex) partialTools[idx] = &partialTool{ id: aws.ToString(toolStart.Value.ToolUseId), name: aws.ToString(toolStart.Value.Name), } } } case *types.ConverseStreamOutputMemberContentBlockStop: // Handle content block stop (tool completion) idx := aws.ToInt32(v.Value.ContentBlockIndex) if partial, exists := partialTools[idx]; exists { // Parse the JSON to extract arguments for function call inputJSON := partial.input.String() var args map[string]any if inputJSON != "" { if err := json.Unmarshal([]byte(inputJSON), &args); err != nil { args = make(map[string]any) } } else { args = make(map[string]any) } // Create ToolUseBlock for conversation history // Use the accumulated JSON string to create proper Input document toolUse := types.ToolUseBlock{ ToolUseId: aws.String(partial.id), Name: aws.String(partial.name), Input: document.NewLazyDocument(args), } completedTools = append(completedTools, toolUse) // Yield tool immediately with parsed arguments response := &bedrockStreamResponse{ content: "", model: c.model, done: false, toolUses: []types.ToolUseBlock{toolUse}, streamingArgs: map[int]map[string]any{0: args}, } if !yield(response, nil) { return } delete(partialTools, idx) } case *types.ConverseStreamOutputMemberMetadata: // Handle final usage metadata if v.Value.Usage != nil { finalResponse := &bedrockStreamResponse{ content: "", usage: v.Value.Usage, model: c.model, done: true, } yield(finalResponse, nil) } } } // Update conversation history with the full response if fullContent.Len() > 0 { assistantMessage.Content = append(assistantMessage.Content, &types.ContentBlockMemberText{Value: fullContent.String()}) } // Include completed tools in conversation history for _, tool := range completedTools { assistantMessage.Content = append(assistantMessage.Content, &types.ContentBlockMemberToolUse{Value: tool}) } // Only add to history if there's content or tools if len(assistantMessage.Content) > 0 { c.messages = append(c.messages, assistantMessage) } // Check for stream errors if err := stream.Err(); err != nil { yield(nil, fmt.Errorf("stream error: %w", err)) } }, nil } // addContentsToHistory processes and appends user messages to chat history // following AWS Bedrock Converse API patterns func (c *bedrockChat) addContentsToHistory(contents []any) error { var contentBlocks []types.ContentBlock for _, content := range contents { switch c := content.(type) { case string: // Add text content block contentBlocks = append(contentBlocks, &types.ContentBlockMemberText{Value: c}) case FunctionCallResult: // Determine status based on Result content status := types.ToolResultStatusSuccess if c.Result != nil { // Check for error field if errorVal, hasError := c.Result["error"]; hasError { if errorBool, isBool := errorVal.(bool); isBool && errorBool { status = types.ToolResultStatusError } } // Check for status field if statusVal, hasStatus := c.Result["status"]; hasStatus { if statusStr, isString := statusVal.(string); isString && (statusStr == "failed" || statusStr == "error") { status = types.ToolResultStatusError } } } // Convert to AWS Bedrock ToolResultBlock format per official docs toolResult := types.ToolResultBlock{ ToolUseId: aws.String(c.ID), Content: []types.ToolResultContentBlock{ &types.ToolResultContentBlockMemberJson{ Value: document.NewLazyDocument(c.Result), }, }, Status: status, } contentBlocks = append(contentBlocks, &types.ContentBlockMemberToolResult{Value: toolResult}) default: return fmt.Errorf("unhandled content type: %T", content) } } if len(contentBlocks) > 0 { // Add user message with all content blocks to conversation history c.messages = append(c.messages, types.Message{ Role: types.ConversationRoleUser, Content: contentBlocks, }) } return nil } // SetFunctionDefinitions configures the available functions for tool use func (c *bedrockChat) SetFunctionDefinitions(functions []*FunctionDefinition) error { c.functionDefs = functions if len(functions) == 0 { c.toolConfig = nil return nil } var tools []types.Tool for _, fn := range functions { // Convert gollm function definition to AWS tool specification inputSchema := make(map[string]interface{}) if fn.Parameters != nil { // Convert Schema to map[string]interface{} jsonData, err := json.Marshal(fn.Parameters) if err != nil { return fmt.Errorf("failed to marshal function parameters: %w", err) } if err := json.Unmarshal(jsonData, &inputSchema); err != nil { return fmt.Errorf("failed to unmarshal function parameters: %w", err) } } toolSpec := types.ToolSpecification{ Name: aws.String(fn.Name), Description: aws.String(fn.Description), InputSchema: &types.ToolInputSchemaMemberJson{ Value: document.NewLazyDocument(inputSchema), }, } tools = append(tools, &types.ToolMemberToolSpec{Value: toolSpec}) } c.toolConfig = &types.ToolConfiguration{ Tools: tools, ToolChoice: &types.ToolChoiceMemberAny{ Value: types.AnyToolChoice{}, }, } return nil } // IsRetryableError determines if an error is retryable func (c *bedrockChat) IsRetryableError(err error) bool { return DefaultIsRetryableError(err) } // bedrockResponse implements ChatResponse for regular (non-streaming) responses type bedrockResponse struct { output *bedrockruntime.ConverseOutput model string } // UsageMetadata returns the usage metadata from the response func (r *bedrockResponse) UsageMetadata() any { if r.output != nil && r.output.Usage != nil { return r.output.Usage } return nil } // Candidates returns the candidate responses func (r *bedrockResponse) Candidates() []Candidate { if r.output == nil || r.output.Output == nil { return []Candidate{} } if msg, ok := r.output.Output.(*types.ConverseOutputMemberMessage); ok { candidate := &bedrockCandidate{ message: &msg.Value, model: r.model, } return []Candidate{candidate} } return []Candidate{} } // bedrockStreamResponse implements ChatResponse for streaming responses type bedrockStreamResponse struct { content string usage *types.TokenUsage model string done bool toolUses []types.ToolUseBlock streamingArgs map[int]map[string]any } // UsageMetadata returns the usage metadata from the streaming response func (r *bedrockStreamResponse) UsageMetadata() any { return r.usage } // Candidates returns the candidate responses for streaming func (r *bedrockStreamResponse) Candidates() []Candidate { if r.content == "" && r.usage == nil && len(r.toolUses) == 0 { return []Candidate{} } candidate := &bedrockStreamCandidate{ content: r.content, model: r.model, toolUses: r.toolUses, streamingArgs: r.streamingArgs, } return []Candidate{candidate} } // bedrockCandidate implements Candidate for regular responses type bedrockCandidate struct { message *types.Message model string } // String returns a string representation of the candidate func (c *bedrockCandidate) String() string { if c.message == nil { return "" } var content strings.Builder for _, block := range c.message.Content { if textBlock, ok := block.(*types.ContentBlockMemberText); ok { content.WriteString(textBlock.Value) } } return content.String() } // Parts returns the parts of the candidate func (c *bedrockCandidate) Parts() []Part { if c.message == nil { return []Part{} } var parts []Part for _, block := range c.message.Content { switch v := block.(type) { case *types.ContentBlockMemberText: parts = append(parts, &bedrockTextPart{text: v.Value}) case *types.ContentBlockMemberToolUse: parts = append(parts, &bedrockToolPart{toolUse: &v.Value}) } } return parts } // bedrockStreamCandidate implements Candidate for streaming responses type bedrockStreamCandidate struct { content string model string toolUses []types.ToolUseBlock streamingArgs map[int]map[string]any } // String returns a string representation of the streaming candidate func (c *bedrockStreamCandidate) String() string { return c.content } // Parts returns the parts of the streaming candidate func (c *bedrockStreamCandidate) Parts() []Part { var parts []Part // Handle text content if c.content != "" { parts = append(parts, &bedrockTextPart{text: c.content}) } // Handle tool calls with streaming args for i, toolUse := range c.toolUses { var args map[string]any if c.streamingArgs != nil { args = c.streamingArgs[i] } parts = append(parts, &bedrockToolPart{ toolUse: &toolUse, args: args, }) } return parts } // bedrockTextPart implements Part for text content type bedrockTextPart struct { text string } // AsText returns the text content func (p *bedrockTextPart) AsText() (string, bool) { return p.text, true } // AsFunctionCalls returns nil since this is a text part func (p *bedrockTextPart) AsFunctionCalls() ([]FunctionCall, bool) { return nil, false } // bedrockToolPart implements Part for tool/function calls type bedrockToolPart struct { toolUse *types.ToolUseBlock args map[string]any // For streaming case when Input can't be unmarshaled } // AsText returns empty string since this is a tool part func (p *bedrockToolPart) AsText() (string, bool) { return "", false } // AsFunctionCalls returns the function calls func (p *bedrockToolPart) AsFunctionCalls() ([]FunctionCall, bool) { if p.toolUse == nil { return nil, false } // Get arguments - prefer pre-parsed args (streaming), fall back to unmarshaling var args map[string]any if p.args != nil { // Streaming case - use pre-parsed arguments args = p.args } else if p.toolUse.Input != nil { // Non-streaming case - unmarshal from Input if err := p.toolUse.Input.UnmarshalSmithyDocument(&args); err != nil { klog.V(2).Infof("Failed to unmarshal tool input: %v", err) args = make(map[string]any) } } else { args = make(map[string]any) } funcCall := FunctionCall{ ID: aws.ToString(p.toolUse.ToolUseId), Name: aws.ToString(p.toolUse.Name), Arguments: args, } return []FunctionCall{funcCall}, true } // Helper functions // getBedrockModel returns the model to use, checking in order: // 1. Explicitly provided model // 2. Environment variable BEDROCK_MODEL // 3. Default model (Claude Sonnet 4) func getBedrockModel(model string) string { if model != "" { klog.V(2).Infof("Using explicitly provided model: %s", model) return model } if envModel := os.Getenv("BEDROCK_MODEL"); envModel != "" { klog.V(1).Infof("Using model from environment variable: %s", envModel) return envModel } defaultModel := "us.anthropic.claude-sonnet-4-20250514-v1:0" klog.V(1).Infof("Using default model: %s", defaultModel) return defaultModel } // bedrockCompletionResponse wraps a ChatResponse to implement CompletionResponse type bedrockCompletionResponse struct { chatResponse ChatResponse } var _ CompletionResponse = (*bedrockCompletionResponse)(nil) func (r *bedrockCompletionResponse) Response() string { if r.chatResponse == nil { return "" } candidates := r.chatResponse.Candidates() if len(candidates) == 0 { return "" } parts := candidates[0].Parts() for _, part := range parts { if text, ok := part.AsText(); ok { return text } } return "" } func (r *bedrockCompletionResponse) UsageMetadata() any { if r.chatResponse == nil { return nil } return r.chatResponse.UsageMetadata() } ================================================ FILE: gollm/factory.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "crypto/tls" "errors" "fmt" "maps" "math/rand/v2" "net" "net/http" "net/url" "os" "slices" "strings" "sync" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "k8s.io/klog/v2" ) var globalRegistry registry type registry struct { mutex sync.Mutex providers map[string]FactoryFunc } func (r *registry) listProviders() []string { r.mutex.Lock() defer r.mutex.Unlock() providers := make([]string, 0, len(r.providers)) for k := range r.providers { providers = append(providers, k) } return providers } type ClientOptions struct { URL *url.URL SkipVerifySSL bool // Extend with more options as needed } // Option is a functional option for configuring ClientOptions. type Option func(*ClientOptions) // WithSkipVerifySSL enables skipping SSL certificate verification for HTTP clients. func WithSkipVerifySSL() Option { return func(o *ClientOptions) { o.SkipVerifySSL = true } } type FactoryFunc func(ctx context.Context, opts ClientOptions) (Client, error) func RegisterProvider(id string, factoryFunc FactoryFunc) error { return globalRegistry.RegisterProvider(id, factoryFunc) } func (r *registry) RegisterProvider(id string, factoryFunc FactoryFunc) error { r.mutex.Lock() defer r.mutex.Unlock() if r.providers == nil { r.providers = make(map[string]FactoryFunc) } _, exists := r.providers[id] if exists { return fmt.Errorf("provider %q is already registered", id) } r.providers[id] = factoryFunc return nil } func (r *registry) NewClient(ctx context.Context, providerID string, opts ...Option) (Client, error) { // providerID can be just an ID, for example "gemini" instead of "gemini://" if !strings.Contains(providerID, "/") && !strings.Contains(providerID, ":") { providerID = providerID + "://" } r.mutex.Lock() defer r.mutex.Unlock() u, err := url.Parse(providerID) if err != nil { return nil, fmt.Errorf("parsing provider id %q: %w", providerID, err) } factoryFunc := r.providers[u.Scheme] if factoryFunc == nil { keys := strings.Join(slices.Collect(maps.Keys(r.providers)), ", ") return nil, fmt.Errorf("provider %q not registered. Available providers: %v", u.Scheme, keys) } // Build ClientOptions clientOpts := ClientOptions{ URL: u, } // Support environment variable override for SkipVerifySSL if v := os.Getenv("LLM_SKIP_VERIFY_SSL"); v == "1" || strings.ToLower(v) == "true" { clientOpts.SkipVerifySSL = true } for _, opt := range opts { opt(&clientOpts) } return factoryFunc(ctx, clientOpts) } /* NewClient builds a Client based on the LLM_CLIENT environment variable or the provided providerID. If providerID is not empty, it overrides the value from LLM_CLIENT. Supports Option parameters and the LLM_SKIP_VERIFY_SSL environment variable. */ func NewClient(ctx context.Context, providerID string, opts ...Option) (Client, error) { if providerID == "" { s := os.Getenv("LLM_CLIENT") if s == "" { return nil, fmt.Errorf("LLM_CLIENT is not set. Available providers: %v", globalRegistry.listProviders()) } providerID = s } return globalRegistry.NewClient(ctx, providerID, opts...) } // APIError represents an error returned by the LLM client. type APIError struct { StatusCode int Message string Err error } func (e *APIError) Error() string { if e.Err != nil { return fmt.Sprintf("API Error: Status=%d, Message='%s', OriginalErr=%v", e.StatusCode, e.Message, e.Err) } return fmt.Sprintf("API Error: Status=%d, Message='%s'", e.StatusCode, e.Message) } func (e *APIError) Unwrap() error { return e.Err } // IsRetryableFunc defines the signature for functions that check if an error is retryable. // TODO (droot): Adjust the signature to allow underlying client to relay the backoff // delay etc. for example, Gemini's error codes contain retryDelay information. type IsRetryableFunc func(error) bool // DefaultIsRetryableError provides a default implementation based on common HTTP codes and network errors. func DefaultIsRetryableError(err error) bool { if err == nil { return false } var apiErr *APIError if errors.As(err, &apiErr) { switch apiErr.StatusCode { case http.StatusConflict, http.StatusTooManyRequests, http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: return true default: return false } } var netErr net.Error if errors.As(err, &netErr) && netErr.Timeout() { return true } // Add other error checks specific to LLM clients if needed // e.g., if errors.Is(err, specificLLMRateLimitError) { return true } return false } // createCustomHTTPClient returns an *http.Client that optionally skips SSL certificate verification. // This is shared by all providers that need custom HTTP transport. func createCustomHTTPClient(skipVerify bool) *http.Client { transport := http.DefaultTransport.(*http.Transport).Clone() transport.Proxy = http.ProxyFromEnvironment if skipVerify { transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} } return &http.Client{ Transport: transport, Timeout: 180 * time.Second, } } // RetryConfig holds the configuration for the retry mechanism (same as before) type RetryConfig struct { MaxAttempts int InitialBackoff time.Duration MaxBackoff time.Duration BackoffFactor float64 Jitter bool } // DefaultRetryConfig provides sensible defaults (same as before) var DefaultRetryConfig = RetryConfig{ MaxAttempts: 5, InitialBackoff: 200 * time.Millisecond, // Slightly increased default MaxBackoff: 10 * time.Second, BackoffFactor: 2.0, Jitter: true, } // Retry executes the provided operation with retries, returning the result and error. // It's now generic to handle any return type T. func Retry[T any]( ctx context.Context, config RetryConfig, isRetryable IsRetryableFunc, operation func(ctx context.Context) (T, error), ) (T, error) { var lastErr error var zero T // Zero value of the return type T log := klog.FromContext(ctx) backoff := config.InitialBackoff for attempt := 1; attempt <= config.MaxAttempts; attempt++ { log.V(2).Info("Retry attempt started", "attempt", attempt, "maxAttempts", config.MaxAttempts, "backoff", backoff) result, err := operation(ctx) if err == nil { log.V(2).Info("Retry attempt succeeded", "attempt", attempt) return result, nil } lastErr = err // Store the last error encountered // Check if context was cancelled *after* the operation select { case <-ctx.Done(): log.Info("Context cancelled after attempt %d failed.", "attempt", attempt) return zero, ctx.Err() // Return context error preferentially default: // Context not cancelled, proceed with error checking } if !isRetryable(lastErr) { log.Info("Attempt failed with non-retryable error", "attempt", attempt, "error", lastErr) return zero, lastErr // Return the non-retryable error immediately } log.Info("Attempt failed with retryable error", "attempt", attempt, "error", lastErr) if attempt == config.MaxAttempts { // Max attempts reached break } // Calculate wait time waitTime := backoff if config.Jitter { waitTime += time.Duration(rand.Float64() * float64(backoff) / 2) } log.V(2).Info("Waiting before next retry attempt", "waitTime", waitTime, "nextAttempt", attempt+1, "maxAttempts", config.MaxAttempts) // Wait or react to context cancellation select { case <-time.After(waitTime): // Wait finished case <-ctx.Done(): log.Info("Context cancelled while waiting for retry after attempt %d.", "attempt", attempt) return zero, ctx.Err() } // Increase backoff backoff = time.Duration(float64(backoff) * config.BackoffFactor) if backoff > config.MaxBackoff { backoff = config.MaxBackoff } } // If the loop finished, it means all attempts failed errFinal := fmt.Errorf("operation failed after %d attempts: %w", config.MaxAttempts, lastErr) return zero, errFinal } // retryChat is a generic decorator that adds retry logic to any Chat implementation. type retryChat[C Chat] struct { underlying Chat // The actual client implementation being wrapped config RetryConfig isRetryable IsRetryableFunc } // NewRetryChat creates a new Chat that wraps the given underlying client // with retry logic using the provided configuration. // It returns the Chat interface type, hiding the generic implementation detail. func NewRetryChat[C Chat]( underlying C, config RetryConfig, ) Chat { return &retryChat[C]{ underlying: underlying, config: config, } } // Embed implements the Client interface for the retryClient decorator. func (rc *retryChat[C]) Send(ctx context.Context, contents ...any) (ChatResponse, error) { // Define the operation operation := func(ctx context.Context) (ChatResponse, error) { return rc.underlying.Send(ctx, contents...) } // Execute with retry return Retry[ChatResponse](ctx, rc.config, rc.underlying.IsRetryableError, operation) } // Embed implements the Client interface for the retryClient decorator. func (rc *retryChat[C]) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { return rc.underlying.SendStreaming(ctx, contents...) } func (rc *retryChat[C]) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error { return rc.underlying.SetFunctionDefinitions(functionDefinitions) } func (rc *retryChat[C]) IsRetryableError(err error) bool { return rc.underlying.IsRetryableError(err) } func (rc *retryChat[C]) Initialize(messages []*api.Message) error { return rc.underlying.Initialize(messages) } ================================================ FILE: gollm/factory_test.go ================================================ // Copyright 2026 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "strings" "testing" ) func TestNewClient(t *testing.T) { _, err := NewClient(context.Background(), "gemini") if err == nil || err.Error() != "GEMINI_API_KEY environment variable not set" { t.Fatalf("Unexpected error: %v", err) } _, err = NewClient(context.Background(), "invalid") if err == nil || !strings.Contains(err.Error(), "provider \"invalid\" not registered") { t.Fatalf("Unexpected error: %v", err) } } ================================================ FILE: gollm/gemini.go ================================================ // Copyright 2024 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "bytes" "context" "encoding/json" "errors" "fmt" "iter" "net" "net/http" "os" "os/exec" "strings" "google.golang.org/genai" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "k8s.io/klog/v2" ) func init() { if err := RegisterProvider("gemini", geminiFactory); err != nil { klog.Fatalf("Failed to register gemini provider: %v", err) } if err := RegisterProvider("vertexai", vertexaiViaGeminiFactory); err != nil { klog.Fatalf("Failed to register vertexai provider: %v", err) } } // geminiFactory is the provider factory function for Gemini. // Supports ClientOptions for consistency, but skipVerifySSL is not used. func geminiFactory(ctx context.Context, opts ClientOptions) (Client, error) { opt := GeminiAPIClientOptions{} return NewGeminiAPIClient(ctx, opt) } // GeminiAPIClientOptions are the options for the Gemini API client. type GeminiAPIClientOptions struct { // API Key for GenAI. Required for BackendGeminiAPI. APIKey string } // NewGeminiAPIClient builds a client for the Gemini API. func NewGeminiAPIClient(ctx context.Context, opt GeminiAPIClientOptions) (*GoogleAIClient, error) { apiKey := opt.APIKey if apiKey == "" { apiKey = os.Getenv("GEMINI_API_KEY") } if apiKey == "" { return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set") } skipVerifySSL := false httpClient := createCustomHTTPClient(skipVerifySSL) httpClient = withJournaling(httpClient) cc := &genai.ClientConfig{ APIKey: apiKey, Backend: genai.BackendGeminiAPI, HTTPClient: httpClient, } client, err := genai.NewClient(ctx, cc) if err != nil { return nil, fmt.Errorf("building gemini client: %w", err) } return &GoogleAIClient{ client: client, }, nil } // VertexAIClientOptions are the options for using the VertexAPI. type VertexAIClientOptions struct { // GCP Project ID for Vertex AI. Required for BackendVertexAI. Project string // GCP Location/Region for Vertex AI. Required for BackendVertexAI. See https://cloud.google.com/vertex-ai/docs/general/locations Location string } // vertexaiViaGeminiFactory is the provider factory function for VertexAI via Gemini. // Supports ClientOptions for consistency, but skipVerifySSL is not used. func vertexaiViaGeminiFactory(ctx context.Context, opts ClientOptions) (Client, error) { opt := VertexAIClientOptions{} return NewVertexAIClient(ctx, opt) } // findDefaultGCPProject gets the default GCP project ID from gcloud func findDefaultGCPProject(ctx context.Context) (string, error) { log := klog.FromContext(ctx) // First check env vars // GOOGLE_CLOUD_PROJECT is the default for the genai library and a GCP convention projectID := "" for _, env := range []string{"GOOGLE_CLOUD_PROJECT"} { if v := os.Getenv(env); v != "" { projectID = v log.Info("got project for vertex client from env var", "project", projectID, "env", env) return projectID, nil } } // Now check default project in gcloud { cmd := exec.CommandContext(ctx, "gcloud", "config", "get", "project") var stdout bytes.Buffer cmd.Stdout = &stdout if err := cmd.Run(); err != nil { return "", fmt.Errorf("cannot get project (using gcloud config get project): %w", err) } projectID = strings.TrimSpace(stdout.String()) if projectID != "" { log.Info("got project from gcloud config", "project", projectID) return projectID, nil } } return "", fmt.Errorf("project was not set in gcloud config (or GOOGLE_CLOUD_PROJECT env var)") } // NewVertexAIClient builds a client for the vertexai API. func NewVertexAIClient(ctx context.Context, opt VertexAIClientOptions) (*GoogleAIClient, error) { log := klog.FromContext(ctx) cc := &genai.ClientConfig{ // Project ID is loaded from the GOOGLE_CLOUD_PROJECT environment variable // Location/Region is loaded from either GOOGLE_CLOUD_LOCATION or GOOGLE_CLOUD_REGION environment variable Backend: genai.BackendVertexAI, Project: opt.Project, Location: opt.Location, } // ProjectID is required if cc.Project == "" { projectID, err := findDefaultGCPProject(ctx) if err != nil { return nil, fmt.Errorf("finding default GCP project ID: %w", err) } cc.Project = projectID } // Location is also required if cc.Location == "" { location := "" // Check well-known env vars for _, env := range []string{"GOOGLE_CLOUD_LOCATION", "GOOGLE_CLOUD_REGION"} { if v := os.Getenv(env); v != "" { location = v log.Info("got location for vertex client from env var", "location", location, "env", env) break } } // Fallback to us-central1 if location == "" { location = "us-central1" log.Info("defaulted location for vertex client", "location", opt.Location) } cc.Location = location } client, err := genai.NewClient(ctx, cc) if err != nil { return nil, fmt.Errorf("building vertexai client: %w", err) } return &GoogleAIClient{ client: client, }, nil } // GoogleAIClient is a client for the google AI APIs. // It implements the Client interface. type GoogleAIClient struct { client *genai.Client // responseSchema will constrain the output to match the given schema responseSchema *genai.Schema } var _ Client = &GoogleAIClient{} // ListModels lists the models available in the Gemini API. func (c *GoogleAIClient) ListModels(ctx context.Context) (modelNames []string, err error) { for model, err := range c.client.Models.All(ctx) { if err != nil { return nil, fmt.Errorf("error listing models: %w", err) } modelNames = append(modelNames, strings.TrimPrefix(model.Name, "models/")) } return modelNames, nil } // Close frees the resources used by the client. func (c *GoogleAIClient) Close() error { return nil } // SetResponseSchema constrains LLM responses to match the provided schema. // Calling with nil will clear the current schema. func (c *GoogleAIClient) SetResponseSchema(responseSchema *Schema) error { if responseSchema == nil { c.responseSchema = nil return nil } geminiSchema, err := toGeminiSchema(responseSchema) if err != nil { return err } c.responseSchema = geminiSchema return nil } func (c *GoogleAIClient) GenerateCompletion(ctx context.Context, request *CompletionRequest) (CompletionResponse, error) { log := klog.FromContext(ctx) var config *genai.GenerateContentConfig if c.responseSchema != nil { config = &genai.GenerateContentConfig{ ResponseSchema: c.responseSchema, ResponseMIMEType: "application/json", } } content := []*genai.Content{ {Role: "user", Parts: []*genai.Part{{Text: request.Prompt}}}, } log.Info("sending GenerateContent request to gemini", "content", content) result, err := c.client.Models.GenerateContent(ctx, request.Model, content, config) if err != nil { return nil, err } return &GeminiCompletionResponse{geminiResponse: result, text: result.Text()}, nil } // StartChat starts a new chat with the model. func (c *GoogleAIClient) StartChat(systemPrompt string, model string) Chat { // Some values that are recommended by aistudio temperature := float32(1.0) topK := float32(40) topP := float32(0.95) maxOutputTokens := int32(8192) chat := &GeminiChat{ model: model, client: c.client, genConfig: &genai.GenerateContentConfig{ SystemInstruction: &genai.Content{ Parts: []*genai.Part{ {Text: systemPrompt}, }, }, Temperature: &temperature, TopK: &topK, TopP: &topP, MaxOutputTokens: maxOutputTokens, ResponseMIMEType: "text/plain", }, history: []*genai.Content{}, } if chat.model == "gemma-3-27b-it" { // Note: gemma-3-27b-it does not allow system prompt // xref: https://discuss.ai.google.dev/t/gemma-3-missing-features-despite-announcement/71692 // TODO: remove this hack once gemma-3-27b-it supports system prompt chat.genConfig.SystemInstruction = nil chat.history = []*genai.Content{ {Role: "user", Parts: []*genai.Part{{Text: systemPrompt}}}, } } if c.responseSchema != nil { chat.genConfig.ResponseSchema = c.responseSchema chat.genConfig.ResponseMIMEType = "application/json" } return chat } // GeminiChat is a chat with the model. // It implements the Chat interface. type GeminiChat struct { model string client *genai.Client history []*genai.Content genConfig *genai.GenerateContentConfig } // SetFunctionDefinitions sets the function definitions for the chat. // This allows the LLM to call user-defined functions. func (c *GeminiChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error { var genaiFunctionDeclarations []*genai.FunctionDeclaration for _, functionDefinition := range functionDefinitions { if functionDefinition.Parameters == nil { return fmt.Errorf("function %q has no parameters", functionDefinition.Name) } parameters, err := toGeminiSchema(functionDefinition.Parameters) if err != nil { return err } genaiFunctionDeclarations = append(genaiFunctionDeclarations, &genai.FunctionDeclaration{ Name: functionDefinition.Name, Description: functionDefinition.Description, Parameters: parameters, }) } c.genConfig.Tools = []*genai.Tool{ { FunctionDeclarations: genaiFunctionDeclarations, }, } return nil } // toGeminiSchema converts our generic Schema to a genai.Schema func toGeminiSchema(schema *Schema) (*genai.Schema, error) { ret := &genai.Schema{ Description: schema.Description, Required: schema.Required, } switch schema.Type { case TypeObject: ret.Type = genai.TypeObject case TypeString: ret.Type = genai.TypeString case TypeNumber: ret.Type = genai.TypeNumber case TypeBoolean: ret.Type = genai.TypeBoolean case TypeInteger: ret.Type = genai.TypeInteger case TypeArray: ret.Type = genai.TypeArray default: return nil, fmt.Errorf("type %q not handled by genai.Schema", schema.Type) } if schema.Properties != nil { ret.Properties = make(map[string]*genai.Schema) for k, v := range schema.Properties { geminiValue, err := toGeminiSchema(v) if err != nil { return nil, err } ret.Properties[k] = geminiValue } } if schema.Items != nil { geminiValue, err := toGeminiSchema(schema.Items) if err != nil { return nil, err } ret.Items = geminiValue } return ret, nil } func (c *GeminiChat) partsToGemini(contents ...any) ([]*genai.Part, error) { var parts []*genai.Part for _, content := range contents { switch v := content.(type) { case string: parts = append(parts, genai.NewPartFromText(v)) case FunctionCallResult: parts = append(parts, &genai.Part{ FunctionResponse: &genai.FunctionResponse{ ID: v.ID, Name: v.Name, Response: v.Result, }, }) default: return nil, fmt.Errorf("unexpected type of content: %T", content) } } return parts, nil } // Send sends a message to the model. // It returns a ChatResponse object containing the response from the model. func (c *GeminiChat) Send(ctx context.Context, contents ...any) (ChatResponse, error) { log := klog.FromContext(ctx) log.V(1).Info("sending LLM request", "user", contents) parts, err := c.partsToGemini(contents...) if err != nil { return nil, err } genaiContent := &genai.Content{ Role: "user", Parts: parts, } c.history = append(c.history, genaiContent) result, err := c.client.Models.GenerateContent(ctx, c.model, c.history, c.genConfig) if err != nil { return nil, fmt.Errorf("failed to generate content: %w", err) } if result == nil || len(result.Candidates) == 0 { return nil, fmt.Errorf("no response from Gemini") } c.history = append(c.history, result.Candidates[0].Content) geminiResponse := result log.V(1).Info("got LLM response", "response", geminiResponse) return &GeminiChatResponse{geminiResponse: geminiResponse}, nil } func (c *GeminiChat) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { log := klog.FromContext(ctx) log.V(1).Info("sending LLM streaming request", "user", contents) parts, err := c.partsToGemini(contents...) if err != nil { return nil, err } genaiContent := &genai.Content{ Role: "user", Parts: parts, } c.history = append(c.history, genaiContent) stream := c.client.Models.GenerateContentStream(ctx, c.model, c.history, c.genConfig) return func(yield func(ChatResponse, error) bool) { next, stop := iter.Pull2(stream) defer stop() for { geminiResponse, err, ok := next() if !ok { return } if err != nil { // Always check for and yield an error first. yield(nil, err) return } if geminiResponse == nil || len(geminiResponse.Candidates) == 0 { return } content := geminiResponse.Candidates[0].Content partsIsEmpty := true if content != nil { for _, part := range content.Parts { if part.Text != "" || part.FunctionCall != nil { partsIsEmpty = false break } } } if partsIsEmpty { // This happens when there is empty content with the finish reason (STOP) to indicate that streaming response is finished. // xref: https://github.com/GoogleCloudPlatform/kubectl-ai/issues/306 log.V(1).Info("empty response probably with STOP finishedReason") return } c.history = append(c.history, content) // yield only when we have a non-empty response if !yield(&GeminiChatResponse{geminiResponse: geminiResponse}, err) { return } } }, nil } func (c *GeminiChat) Initialize(messages []*api.Message) error { klog.Info("Initializing gemini chat") c.history = make([]*genai.Content, 0, len(messages)) for _, msg := range messages { content, err := c.messageToContent(msg) if err != nil { continue } c.history = append(c.history, content) } return nil } func (c *GeminiChat) messageToContent(msg *api.Message) (*genai.Content, error) { var role string switch msg.Source { case api.MessageSourceUser: role = "user" case api.MessageSourceModel: role = "model" case api.MessageSourceAgent: role = "user" // Treat agent messages as user messages for Gemini history default: return nil, fmt.Errorf("unknown message source: %s", msg.Source) } parts, err := c.partsToGemini(msg.Payload) if err != nil { return nil, fmt.Errorf("failed to convert message payload to parts: %w", err) } return &genai.Content{Role: role, Parts: parts}, nil } // GeminiChatResponse is a response from the Gemini API. // It implements the ChatResponse interface. type GeminiChatResponse struct { geminiResponse *genai.GenerateContentResponse } var _ ChatResponse = &GeminiChatResponse{} func (r *GeminiChatResponse) MarshalJSON() ([]byte, error) { formatted := RecordChatResponse{ Raw: r.geminiResponse, } return json.Marshal(&formatted) } // String returns a string representation of the response. func (r *GeminiChatResponse) String() string { return r.geminiResponse.Text() } // UsageMetadata returns the usage metadata for the response. func (r *GeminiChatResponse) UsageMetadata() any { return r.geminiResponse.UsageMetadata } // Candidates returns the candidates for the response. func (r *GeminiChatResponse) Candidates() []Candidate { var candidates []Candidate for _, candidate := range r.geminiResponse.Candidates { candidates = append(candidates, &GeminiCandidate{candidate: candidate}) } return candidates } // GeminiCandidate is a candidate for the response. // It implements the Candidate interface. type GeminiCandidate struct { candidate *genai.Candidate } // String returns a string representation of the response. func (r *GeminiCandidate) String() string { var response strings.Builder response.WriteString("[") for i, parts := range r.Parts() { if i > 0 { response.WriteString(", ") } text, ok := parts.AsText() if ok { response.WriteString(text) } functionCalls, ok := parts.AsFunctionCalls() if ok { response.WriteString("functionCalls=[") for _, functionCall := range functionCalls { response.WriteString(fmt.Sprintf("%q(args=%v)", functionCall.Name, functionCall.Arguments)) } response.WriteString("]}") } } response.WriteString("]}") return response.String() } // Parts returns the parts of the candidate. func (r *GeminiCandidate) Parts() []Part { var parts []Part if r.candidate.Content != nil { for _, part := range r.candidate.Content.Parts { parts = append(parts, &GeminiPart{part: *part}) } } return parts } // GeminiPart is a part of a candidate. // It implements the Part interface. type GeminiPart struct { part genai.Part } // AsText returns the text of the part. func (p *GeminiPart) AsText() (string, bool) { if p.part.Text != "" { return p.part.Text, true } return "", false } // AsFunctionCalls returns the function calls of the part. func (p *GeminiPart) AsFunctionCalls() ([]FunctionCall, bool) { if p.part.FunctionCall != nil { return []FunctionCall{ { ID: p.part.FunctionCall.ID, Name: p.part.FunctionCall.Name, Arguments: p.part.FunctionCall.Args, }, }, true } return nil, false } type GeminiCompletionResponse struct { geminiResponse *genai.GenerateContentResponse text string } var _ CompletionResponse = &GeminiCompletionResponse{} func (r *GeminiCompletionResponse) MarshalJSON() ([]byte, error) { formatted := RecordCompletionResponse{ Text: r.text, Raw: r.geminiResponse, } return json.Marshal(&formatted) } func (r *GeminiCompletionResponse) Response() string { return r.text } func (r *GeminiCompletionResponse) UsageMetadata() any { return r.geminiResponse.UsageMetadata } func (r *GeminiCompletionResponse) String() string { return fmt.Sprintf("{text=%q}", r.text) } func (c *GeminiChat) IsRetryableError(err error) bool { if err == nil { return false } var apiErr genai.APIError if errors.As(err, &apiErr) { switch apiErr.Code { case http.StatusConflict, http.StatusTooManyRequests, http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: return true default: return false } } var netErr net.Error if errors.As(err, &netErr) && netErr.Timeout() { return true } // Add other error checks specific to LLM clients if needed // e.g., if errors.Is(err, specificLLMRateLimitError) { return true } return false } ================================================ FILE: gollm/go.mod ================================================ module github.com/GoogleCloudPlatform/kubectl-ai/gollm go 1.24.0 toolchain go1.24.3 require ( github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.7.2 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.7.0 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.2.0 github.com/GoogleCloudPlatform/kubectl-ai v0.0.19 github.com/anthropics/anthropic-sdk-go v1.26.0 github.com/aws/aws-sdk-go-v2 v1.36.6 github.com/aws/aws-sdk-go-v2/config v1.29.18 github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.31.1 github.com/ollama/ollama v0.6.5 github.com/openai/openai-go v1.11.0 google.golang.org/genai v1.8.0 k8s.io/klog/v2 v2.130.1 ) require ( cloud.google.com/go v0.118.3 // indirect cloud.google.com/go/auth v0.15.0 // indirect cloud.google.com/go/compute/metadata v0.6.0 // indirect github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 // indirect github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.11 // indirect github.com/aws/aws-sdk-go-v2/credentials v1.17.71 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.33 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.37 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.37 // indirect github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.4 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.18 // indirect github.com/aws/aws-sdk-go-v2/service/sso v1.25.6 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.34.1 // indirect github.com/aws/smithy-go v1.22.4 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.2 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/golang-jwt/jwt/v5 v5.2.2 // indirect github.com/google/go-cmp v0.7.0 // indirect github.com/google/s2a-go v0.1.9 // indirect github.com/google/uuid v1.6.0 // indirect github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect github.com/googleapis/gax-go/v2 v2.14.1 // indirect github.com/gorilla/websocket v1.5.3 // indirect github.com/kylelemons/godebug v1.1.0 // indirect github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect go.opentelemetry.io/auto/sdk v1.1.0 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 // indirect go.opentelemetry.io/otel v1.34.0 // indirect go.opentelemetry.io/otel/metric v1.34.0 // indirect go.opentelemetry.io/otel/trace v1.34.0 // indirect golang.org/x/crypto v0.40.0 // indirect golang.org/x/net v0.41.0 // indirect golang.org/x/sync v0.16.0 // indirect golang.org/x/sys v0.34.0 // indirect golang.org/x/text v0.27.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2 // indirect google.golang.org/grpc v1.70.0 // indirect google.golang.org/protobuf v1.36.5 // indirect sigs.k8s.io/yaml v1.4.0 // indirect ) ================================================ FILE: gollm/go.sum ================================================ cloud.google.com/go v0.118.3 h1:jsypSnrE/w4mJysioGdMBg4MiW/hHx/sArFpaBWHdME= cloud.google.com/go v0.118.3/go.mod h1:Lhs3YLnBlwJ4KA6nuObNMZ/fCbOQBPuWKPoE0Wa/9Vc= cloud.google.com/go/auth v0.15.0 h1:Ly0u4aA5vG/fsSsxu98qCQBemXtAtJf+95z9HK+cxps= cloud.google.com/go/auth v0.15.0/go.mod h1:WJDGqZ1o9E9wKIL+IwStfyn/+s59zl4Bi+1KQNVXLZ8= cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I= cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg= github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.7.2 h1:+hDUZnYHHoXu05iXiJcL53MZW7raZZejB8ZtzVW7yyc= github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.7.2/go.mod h1:49PyorVrwk6G+e8Vghvn7EkAS6wSPdXEu5a8iW2/vC8= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 h1:Gt0j3wceWMwPmiazCa8MzMA0MfhmPIz0Qp0FJ6qcM0U= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0/go.mod h1:Ot/6aikWnKWi4l9QB7qVSwa8iMphQNqkWALMoNT3rzM= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0 h1:OVoM452qUFBrX+URdH3VpR299ma4kfom0yB0URYky9g= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0/go.mod h1:kUjrAo8bgEwLeZ/CmHqNl3Z/kPm7y6FKfxxK0izYUg4= github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2 h1:yz1bePFlP5Vws5+8ez6T3HWXPmwOK7Yvq8QxDBD3SKY= github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2/go.mod h1:Pa9ZNPuoNu/GztvBSKk9J1cDJW6vk/n0zLtV4mgd8N8= github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 h1:FPKJS1T+clwv+OLGt13a8UjqeRuh0O4SJ3lUriThc+4= github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1/go.mod h1:j2chePtV91HrC22tGoRX3sGY42uF13WzmmV80/OdVAA= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.7.0 h1:4exaC92+n1FzhSKb5Ghino2XEk3cClUtzvveL1U9YeM= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.7.0/go.mod h1:BkhZrH3JiVTkrTqCeYHOmqReFcZTYEMf8jcFDlrCJLk= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.2.0 h1:UrGzkHueDwAWDdjQxC+QaXHd4tVCkISYE9j7fSSXF8k= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.2.0/go.mod h1:qskvSQeW+cxEE2bcKYyKimB1/KiQ9xpJ99bcHY0BX6c= github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJTmL004Abzc5wDB5VtZG2PJk5ndYDgVacGqfirKxjM= github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1/go.mod h1:tCcJZ0uHAmvjsVYzEFivsRTN00oz5BEsRgQHu5JZ9WE= github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2 h1:oygO0locgZJe7PpYPXT5A29ZkwJaPqcva7BVeemZOZs= github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI= github.com/GoogleCloudPlatform/kubectl-ai v0.0.19 h1:RdVCft8obsRZaoyVjqaOMu+ylnPb+CKcd8pRYPq3zvs= github.com/GoogleCloudPlatform/kubectl-ai v0.0.19/go.mod h1:VOHud1Et2RE668c2dcdxApYclNk74SjKLxoaQQtK6Bc= github.com/anthropics/anthropic-sdk-go v1.26.0 h1:oUTzFaUpAevfuELAP1sjL6CQJ9HHAfT7CoSYSac11PY= github.com/anthropics/anthropic-sdk-go v1.26.0/go.mod h1:qUKmaW+uuPB64iy1l+4kOSvaLqPXnHTTBKH6RVZ7q5Q= github.com/aws/aws-sdk-go-v2 v1.36.6 h1:zJqGjVbRdTPojeCGWn5IR5pbJwSQSBh5RWFTQcEQGdU= github.com/aws/aws-sdk-go-v2 v1.36.6/go.mod h1:EYrzvCCN9CMUTa5+6lf6MM4tq3Zjp8UhSGR/cBsjai0= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.11 h1:12SpdwU8Djs+YGklkinSSlcrPyj3H4VifVsKf78KbwA= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.11/go.mod h1:dd+Lkp6YmMryke+qxW/VnKyhMBDTYP41Q2Bb+6gNZgY= github.com/aws/aws-sdk-go-v2/config v1.29.18 h1:x4T1GRPnqKV8HMJOMtNktbpQMl3bIsfx8KbqmveUO2I= github.com/aws/aws-sdk-go-v2/config v1.29.18/go.mod h1:bvz8oXugIsH8K7HLhBv06vDqnFv3NsGDt2Znpk7zmOU= github.com/aws/aws-sdk-go-v2/credentials v1.17.71 h1:r2w4mQWnrTMJjOyIsZtGp3R3XGY3nqHn8C26C2lQWgA= github.com/aws/aws-sdk-go-v2/credentials v1.17.71/go.mod h1:E7VF3acIup4GB5ckzbKFrCK0vTvEQxOxgdq4U3vcMCY= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.33 h1:D9ixiWSG4lyUBL2DDNK924Px9V/NBVpML90MHqyTADY= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.33/go.mod h1:caS/m4DI+cij2paz3rtProRBI4s/+TCiWoaWZuQ9010= github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.37 h1:osMWfm/sC/L4tvEdQ65Gri5ZZDCUpuYJZbTTDrsn4I0= github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.37/go.mod h1:ZV2/1fbjOPr4G4v38G3Ww5TBT4+hmsK45s/rxu1fGy0= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.37 h1:v+X21AvTb2wZ+ycg1gx+orkB/9U6L7AOp93R7qYxsxM= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.37/go.mod h1:G0uM1kyssELxmJ2VZEfG0q2npObR3BAkF3c1VsfVnfs= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 h1:bIqFDwgGXXN1Kpp99pDOdKMTTb5d2KyU5X/BZxjOkRo= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3/go.mod h1:H5O/EsxDWyU+LP/V8i5sm8cxoZgc2fdNR9bxlOFrQTo= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.31.1 h1:JDLT1baDmioiZKa2bZ6J82/Zwfv9cSAjr+LyF47TPYw= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.31.1/go.mod h1:FvbGcqrU4sC3qjrAKK3FzOmBoucDJF2dXsKVvAbGE8g= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.4 h1:CXV68E2dNqhuynZJPB80bhPQwAKqBWVer887figW6Jc= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.4/go.mod h1:/xFi9KtvBXP97ppCz1TAEvU1Uf66qvid89rbem3wCzQ= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.18 h1:vvbXsA2TVO80/KT7ZqCbx934dt6PY+vQ8hZpUZ/cpYg= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.18/go.mod h1:m2JJHledjBGNMsLOF1g9gbAxprzq3KjC8e4lxtn+eWg= github.com/aws/aws-sdk-go-v2/service/sso v1.25.6 h1:rGtWqkQbPk7Bkwuv3NzpE/scwwL9sC1Ul3tn9x83DUI= github.com/aws/aws-sdk-go-v2/service/sso v1.25.6/go.mod h1:u4ku9OLv4TO4bCPdxf4fA1upaMaJmP9ZijGk3AAOC6Q= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4 h1:OV/pxyXh+eMA0TExHEC4jyWdumLxNbzz1P0zJoezkJc= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4/go.mod h1:8Mm5VGYwtm+r305FfPSuc+aFkrypeylGYhFim6XEPoc= github.com/aws/aws-sdk-go-v2/service/sts v1.34.1 h1:aUrLQwJfZtwv3/ZNG2xRtEen+NqI3iesuacjP51Mv1s= github.com/aws/aws-sdk-go-v2/service/sts v1.34.1/go.mod h1:3wFBZKoWnX3r+Sm7in79i54fBmNfwhdNdQuscCw7QIk= github.com/aws/smithy-go v1.22.4 h1:uqXzVZNuNexwc/xrh6Tb56u89WDlJY6HS+KC0S4QSjw= github.com/aws/smithy-go v1.22.4/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI= github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw= github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA= github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q= github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/keybase/go-keychain v0.0.1 h1:way+bWYa6lDppZoZcgMbYsvC7GxljxrskdNInRtuthU= github.com/keybase/go-keychain v0.0.1/go.mod h1:PdEILRW3i9D8JcdM+FmY6RwkHGnhHxXwkPPMeUgOK1k= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/ollama/ollama v0.6.5 h1:vXKkVX57ql/1ZzMw4SVK866Qfd6pjwEcITVyEpF0QXQ= github.com/ollama/ollama v0.6.5/go.mod h1:pGgtoNyc9DdM6oZI6yMfI6jTk2Eh4c36c2GpfQCH7PY= github.com/openai/openai-go v1.11.0 h1:ztH+W0ug5Kh9+/EErHa8KAmhwixkzjK57rXyE+ZnSCk= github.com/openai/openai-go v1.11.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM= github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 h1:CV7UdSGJt/Ao6Gp4CXckLxVRRsRgDHoI8XjbL3PDl8s= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0/go.mod h1:FRmFuRJfag1IZ2dPkHnEoSFVgTVPUd2qf5Vi69hLb8I= go.opentelemetry.io/otel v1.34.0 h1:zRLXxLCgL1WyKsPVrgbSdMN4c0FMkDAskSTQP+0hdUY= go.opentelemetry.io/otel v1.34.0/go.mod h1:OWFPOQ+h4G8xpyjgqo4SxJYdDQ/qmRH+wivy7zzx9oI= go.opentelemetry.io/otel/metric v1.34.0 h1:+eTR3U0MyfWjRDhmFMxe2SsW64QrZ84AOhvqS7Y+PoQ= go.opentelemetry.io/otel/metric v1.34.0/go.mod h1:CEDrp0fy2D0MvkXE+dPV7cMi8tWZwX3dmaIhwPOaqHE= go.opentelemetry.io/otel/sdk v1.34.0 h1:95zS4k/2GOy069d321O8jWgYsW3MzVV+KuSPKp7Wr1A= go.opentelemetry.io/otel/sdk v1.34.0/go.mod h1:0e/pNiaMAqaykJGKbi+tSjWfNNHMTxoC9qANsCzbyxU= go.opentelemetry.io/otel/sdk/metric v1.32.0 h1:rZvFnvmvawYb0alrYkjraqJq0Z4ZUJAiyYCU9snn1CU= go.opentelemetry.io/otel/sdk/metric v1.32.0/go.mod h1:PWeZlq0zt9YkYAp3gjKZ0eicRYvOh1Gd+X99x6GHpCQ= go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC8mh/k= go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE= golang.org/x/crypto v0.40.0 h1:r4x+VvoG5Fm+eJcxMaY8CQM7Lb0l1lsmjGBQ6s8BfKM= golang.org/x/crypto v0.40.0/go.mod h1:Qr1vMER5WyS2dfPHAlsOj01wgLbsyWtFn/aY+5+ZdxY= golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw= golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA= golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw= golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.34.0 h1:H5Y5sJ2L2JRdyv7ROF1he/lPdvFsd0mJHFw2ThKHxLA= golang.org/x/sys v0.34.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/text v0.27.0 h1:4fGWRpyh641NLlecmyl4LOe6yDdfaYNrGb2zdfo4JV4= golang.org/x/text v0.27.0/go.mod h1:1D28KMCvyooCX9hBiosv5Tz/+YLxj0j7XhWjpSUF7CU= google.golang.org/genai v1.8.0 h1:unX2CNWSiKDO2MSTKK3RstXg/vHp9hr42LIcL6f3Cik= google.golang.org/genai v1.8.0/go.mod h1:TyfOKRz/QyCaj6f/ZDt505x+YreXnY40l2I6k8TvgqY= google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2 h1:DMTIbak9GhdaSxEjvVzAeNZvyc03I61duqNbnm3SU0M= google.golang.org/genproto/googleapis/rpc v0.0.0-20250219182151-9fdb1cabc7b2/go.mod h1:LuRYeWDFV6WOn90g357N17oMCaxpgCnbi/44qJvDn2I= google.golang.org/grpc v1.70.0 h1:pWFv03aZoHzlRKHWicjsZytKAiYCtNS0dHbXnIdq7jQ= google.golang.org/grpc v1.70.0/go.mod h1:ofIJqVKDXx/JiXrwr2IG4/zwdH9txy3IlF40RmcJSQw= google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM= google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk= k8s.io/klog/v2 v2.130.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE= sigs.k8s.io/yaml v1.4.0 h1:Mk1wCc2gy/F0THH0TAp1QYyJNzRm2KCLy3o5ASXVI5E= sigs.k8s.io/yaml v1.4.0/go.mod h1:Ejl7/uTz7PSA4eKMyQCUTnhZYNmLIl+5c2lQPGR2BPY= ================================================ FILE: gollm/grok.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "errors" "fmt" "os" openai "github.com/openai/openai-go" "github.com/openai/openai-go/option" "k8s.io/klog/v2" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) // Register the Grok provider factory on package initialization. // The new factory function supports ClientOptions, including skipVerifySSL. func init() { if err := RegisterProvider("grok", newGrokClientFactory); err != nil { klog.Fatalf("Failed to register Grok provider: %v", err) } } // newGrokClientFactory is the factory function for creating Grok clients with options. func newGrokClientFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewGrokClient(ctx, opts) } // GrokClient implements the gollm.Client interface for X.AI's Grok model. type GrokClient struct { client openai.Client } // Ensure GrokClient implements the Client interface. var _ Client = &GrokClient{} // NewGrokClient creates a new client for interacting with X.AI's Grok model. // Supports custom HTTP client and skipVerifySSL via ClientOptions. func NewGrokClient(ctx context.Context, opts ClientOptions) (*GrokClient, error) { apiKey := os.Getenv("GROK_API_KEY") if apiKey == "" { return nil, errors.New("GROK_API_KEY environment variable not set") } // Default API endpoint for X.AI endpoint := "https://api.x.ai/v1" // Allow endpoint override customEndpoint := os.Getenv("GROK_ENDPOINT") if customEndpoint != "" { endpoint = customEndpoint klog.Infof("Using custom Grok endpoint: %s", endpoint) } // Use the OpenAI client with custom base URL and custom HTTP client httpClient := createCustomHTTPClient(opts.SkipVerifySSL) return &GrokClient{ client: openai.NewClient( option.WithAPIKey(apiKey), option.WithBaseURL(endpoint), option.WithHTTPClient(httpClient), ), }, nil } // Close cleans up any resources used by the client. func (c *GrokClient) Close() error { // No specific cleanup needed for the Grok client currently. return nil } // StartChat starts a new chat session. func (c *GrokClient) StartChat(systemPrompt, model string) Chat { // Default to Grok-3-beta if no model is specified if model == "" { model = "grok-3-beta" klog.V(1).Info("No model specified, defaulting to grok-3-beta") } klog.V(1).Infof("Starting new Grok chat session with model: %s", model) // Initialize history with system prompt if provided history := []openai.ChatCompletionMessageParamUnion{} if systemPrompt != "" { history = append(history, openai.SystemMessage(systemPrompt)) } return &grokChatSession{ client: c.client, history: history, model: model, } } // simpleGrokCompletionResponse is a basic implementation of CompletionResponse. type simpleGrokCompletionResponse struct { content string } // Response returns the completion content. func (r *simpleGrokCompletionResponse) Response() string { return r.content } // UsageMetadata returns nil for now. func (r *simpleGrokCompletionResponse) UsageMetadata() any { return nil } // GenerateCompletion sends a completion request to the Grok API. func (c *GrokClient) GenerateCompletion(ctx context.Context, req *CompletionRequest) (CompletionResponse, error) { klog.Infof("Grok GenerateCompletion called with model: %s", req.Model) klog.V(1).Infof("Prompt:\n%s", req.Prompt) // Use the Chat Completions API as shown in examples chatReq := openai.ChatCompletionNewParams{ Model: openai.ChatModel(req.Model), // Use the model specified in the request Messages: []openai.ChatCompletionMessageParamUnion{ // Assuming a simple user message structure for now openai.UserMessage(req.Prompt), }, } completion, err := c.client.Chat.Completions.New(ctx, chatReq) if err != nil { return nil, fmt.Errorf("failed to generate Grok completion: %w", err) } // Check if there are choices and a message if len(completion.Choices) == 0 || completion.Choices[0].Message.Content == "" { return nil, errors.New("received an empty response from Grok") } // Return the content of the first choice resp := &simpleGrokCompletionResponse{ content: completion.Choices[0].Message.Content, } return resp, nil } // SetResponseSchema is not implemented yet for Grok. func (c *GrokClient) SetResponseSchema(schema *Schema) error { klog.Warning("GrokClient.SetResponseSchema is not implemented yet") return nil } // ListModels returns a list of available Grok models. func (c *GrokClient) ListModels(ctx context.Context) ([]string, error) { // Currently, Grok only has a fixed set of models // This could be updated to call a models endpoint if X.AI provides one in the future return []string{"grok-3-beta"}, nil } // --- Chat Session Implementation --- type grokChatSession struct { client openai.Client history []openai.ChatCompletionMessageParamUnion model string functionDefinitions []*FunctionDefinition // Stored in gollm format tools []openai.ChatCompletionToolParam // Stored in OpenAI format } // Ensure grokChatSession implements the Chat interface. var _ Chat = (*grokChatSession)(nil) // SetFunctionDefinitions stores the function definitions and converts them to Grok format. func (cs *grokChatSession) SetFunctionDefinitions(defs []*FunctionDefinition) error { cs.functionDefinitions = defs cs.tools = nil // Clear previous tools if len(defs) > 0 { cs.tools = make([]openai.ChatCompletionToolParam, len(defs)) for i, gollmDef := range defs { // Basic conversion, assuming schema is compatible or nil var params openai.FunctionParameters if gollmDef.Parameters != nil { // NOTE: This assumes gollmDef.Parameters is directly marshalable to JSON // that fits openai.FunctionParameters. May need refinement. bytes, err := gollmDef.Parameters.ToRawSchema() if err != nil { return fmt.Errorf("failed to convert schema for function %s: %w", gollmDef.Name, err) } if err := json.Unmarshal(bytes, ¶ms); err != nil { return fmt.Errorf("failed to unmarshal schema for function %s: %w", gollmDef.Name, err) } } cs.tools[i] = openai.ChatCompletionToolParam{ Function: openai.FunctionDefinitionParam{ Name: gollmDef.Name, Description: openai.String(gollmDef.Description), Parameters: params, }, } } } klog.V(1).Infof("Set %d function definitions for Grok chat session", len(cs.functionDefinitions)) return nil } // Send sends the user message(s), appends to history, and gets the LLM response. func (cs *grokChatSession) Send(ctx context.Context, contents ...any) (ChatResponse, error) { klog.V(1).InfoS("grokChatSession.Send called", "model", cs.model, "history_len", len(cs.history)) // Append user message(s) to history for _, content := range contents { switch c := content.(type) { case string: klog.V(2).Infof("Adding user message to history: %s", c) cs.history = append(cs.history, openai.UserMessage(c)) case FunctionCallResult: klog.V(2).Infof("Adding tool call result to history: Name=%s, ID=%s", c.Name, c.ID) // Marshal the result map into a JSON string for the message content resultJSON, err := json.Marshal(c.Result) if err != nil { klog.Errorf("Failed to marshal function call result: %v", err) return nil, fmt.Errorf("failed to marshal function call result %q: %w", c.Name, err) } cs.history = append(cs.history, openai.ToolMessage(string(resultJSON), c.ID)) default: // TODO: Handle other content types if necessary? klog.Warningf("Unhandled content type in Send: %T", content) return nil, fmt.Errorf("unhandled content type: %T", content) } } // Prepare the API request chatReq := openai.ChatCompletionNewParams{ Model: openai.ChatModel(cs.model), Messages: cs.history, } if len(cs.tools) > 0 { chatReq.Tools = cs.tools // chatReq.ToolChoice = openai.ToolChoiceAuto // Or specify if needed } // Call the Grok API klog.V(1).InfoS("Sending request to Grok Chat API", "model", cs.model, "messages", len(chatReq.Messages), "tools", len(chatReq.Tools)) completion, err := cs.client.Chat.Completions.New(ctx, chatReq) if err != nil { klog.Errorf("Grok ChatCompletion API error: %v", err) return nil, fmt.Errorf("Grok chat completion failed: %w", err) } klog.V(1).InfoS("Received response from Grok Chat API", "id", completion.ID, "choices", len(completion.Choices)) // Process the response if len(completion.Choices) == 0 { klog.Warning("Received response with no choices from Grok") return nil, errors.New("received empty response from Grok (no choices)") } // Add assistant's response (first choice) to history assistantMsg := completion.Choices[0].Message // Convert to param type before appending to history cs.history = append(cs.history, assistantMsg.ToParam()) klog.V(2).InfoS("Added assistant message to history", "content_present", assistantMsg.Content != "", "tool_calls", len(assistantMsg.ToolCalls)) // Wrap the response resp := &grokChatResponse{ grokCompletion: completion, } return resp, nil } // SendStreaming sends the user message(s) and returns an iterator for the LLM response stream. func (cs *grokChatSession) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { klog.V(1).InfoS("Starting Grok streaming request", "model", cs.model, "streamingEnabled", true) // Append user message(s) to history for _, content := range contents { switch c := content.(type) { case string: klog.V(2).Infof("Adding user message to history: %s", c) cs.history = append(cs.history, openai.UserMessage(c)) case FunctionCallResult: klog.V(2).Infof("Adding tool call result to history: Name=%s, ID=%s", c.Name, c.ID) resultJSON, err := json.Marshal(c.Result) if err != nil { klog.Errorf("Failed to marshal function call result: %v", err) return nil, fmt.Errorf("failed to marshal function call result %q: %w", c.Name, err) } cs.history = append(cs.history, openai.ToolMessage(string(resultJSON), c.ID)) default: klog.Warningf("Unhandled content type in SendStreaming: %T", content) return nil, fmt.Errorf("unhandled content type: %T", content) } } // Prepare the API request chatReq := openai.ChatCompletionNewParams{ Model: openai.ChatModel(cs.model), Messages: cs.history, } if len(cs.tools) > 0 { chatReq.Tools = cs.tools } // Start the Grok streaming request klog.V(1).InfoS("Sending streaming request to Grok API", "model", cs.model, "messageCount", len(chatReq.Messages), "toolCount", len(chatReq.Tools)) stream := cs.client.Chat.Completions.NewStreaming(ctx, chatReq) // Create an accumulator to track the full response acc := openai.ChatCompletionAccumulator{} // Create and return the stream iterator return func(yield func(ChatResponse, error) bool) { var lastResponseChunk *grokChatStreamResponse // Process stream chunks for stream.Next() { chunk := stream.Current() // Update the accumulator with the new chunk acc.AddChunk(chunk) // Create a streaming response for this chunk streamResponse := &grokChatStreamResponse{ streamChunk: chunk, accumulator: acc, } // Keep track of the last response to append to history lastResponseChunk = streamResponse // Yield the streaming response if !yield(streamResponse, nil) { // Consumer wants to stop break } } // Check for errors after streaming completes if err := stream.Err(); err != nil { klog.Errorf("Error in Grok streaming: %v", err) yield(nil, fmt.Errorf("Grok streaming error: %w", err)) return } // Update conversation history with the complete message if lastResponseChunk != nil && acc.Choices != nil && len(acc.Choices) > 0 { // The accumulator has the complete message completeMessage := openai.ChatCompletionMessage{ Content: acc.Choices[0].Message.Content, Role: acc.Choices[0].Message.Role, ToolCalls: acc.Choices[0].Message.ToolCalls, } // Append the full assistant response to history cs.history = append(cs.history, completeMessage.ToParam()) klog.V(2).InfoS("Added complete assistant message to history", "content_present", completeMessage.Content != "", "tool_calls", len(completeMessage.ToolCalls)) } }, nil } // IsRetryableError determines if an error from the Grok API should be retried. func (cs *grokChatSession) IsRetryableError(err error) bool { if err == nil { return false } return DefaultIsRetryableError(err) } func (cs *grokChatSession) Initialize(messages []*api.Message) error { klog.Warning("chat history persistence is not supported for provider 'grok', using in-memory chat history") return nil } // --- Helper structs for ChatResponse interface --- type grokChatResponse struct { grokCompletion *openai.ChatCompletion } var _ ChatResponse = (*grokChatResponse)(nil) func (r *grokChatResponse) UsageMetadata() any { // Check if the main completion object and Usage exist if r.grokCompletion != nil && r.grokCompletion.Usage.TotalTokens > 0 { // Check a field within Usage return r.grokCompletion.Usage } return nil } func (r *grokChatResponse) Candidates() []Candidate { if r.grokCompletion == nil { return nil } candidates := make([]Candidate, len(r.grokCompletion.Choices)) for i, choice := range r.grokCompletion.Choices { candidates[i] = &grokCandidate{grokChoice: &choice} } return candidates } type grokCandidate struct { grokChoice *openai.ChatCompletionChoice } var _ Candidate = (*grokCandidate)(nil) func (c *grokCandidate) Parts() []Part { // Check if the choice exists before accessing Message if c.grokChoice == nil { return nil } // Grok message can have Content AND ToolCalls var parts []Part if c.grokChoice.Message.Content != "" { parts = append(parts, &grokPart{content: c.grokChoice.Message.Content}) } if len(c.grokChoice.Message.ToolCalls) > 0 { parts = append(parts, &grokPart{toolCalls: c.grokChoice.Message.ToolCalls}) } return parts } // String provides a simple string representation for logging/debugging. func (c *grokCandidate) String() string { if c.grokChoice == nil { return "" } content := "" if c.grokChoice.Message.Content != "" { content = c.grokChoice.Message.Content } toolCalls := len(c.grokChoice.Message.ToolCalls) finishReason := string(c.grokChoice.FinishReason) return fmt.Sprintf("Candidate(FinishReason: %s, ToolCalls: %d, Content: %q)", finishReason, toolCalls, content) } type grokPart struct { content string toolCalls []openai.ChatCompletionMessageToolCall } var _ Part = (*grokPart)(nil) func (p *grokPart) AsText() (string, bool) { return p.content, p.content != "" } func (p *grokPart) AsFunctionCalls() ([]FunctionCall, bool) { if len(p.toolCalls) == 0 { return nil, false } gollmCalls := make([]FunctionCall, len(p.toolCalls)) for i, tc := range p.toolCalls { // Check if it's a function call by seeing if Function Name is populated if tc.Function.Name == "" { klog.V(2).Infof("Skipping non-function tool call ID: %s", tc.ID) continue } var args map[string]any // Attempt to unmarshal arguments, ignore error for now if it fails _ = json.Unmarshal([]byte(tc.Function.Arguments), &args) gollmCalls[i] = FunctionCall{ ID: tc.ID, Name: tc.Function.Name, Arguments: args, } } return gollmCalls, true } // grokChatStreamResponse represents a streaming response chunk from Grok. type grokChatStreamResponse struct { streamChunk openai.ChatCompletionChunk accumulator openai.ChatCompletionAccumulator } // Ensure the streaming response implements ChatResponse interface. var _ ChatResponse = (*grokChatStreamResponse)(nil) // UsageMetadata returns usage metadata if available in the final chunk. func (r *grokChatStreamResponse) UsageMetadata() any { if r.accumulator.Usage.TotalTokens > 0 { return r.accumulator.Usage } return nil } // Candidates returns a slice with a single streaming candidate. func (r *grokChatStreamResponse) Candidates() []Candidate { // Each streaming chunk gets converted to a candidate if len(r.streamChunk.Choices) == 0 { return nil } candidates := make([]Candidate, len(r.streamChunk.Choices)) for i, choice := range r.streamChunk.Choices { candidates[i] = &grokStreamCandidate{streamChoice: choice} } return candidates } // grokStreamCandidate adapts a streaming chunk choice to the Candidate interface. type grokStreamCandidate struct { streamChoice openai.ChatCompletionChunkChoice } // Ensure the streaming candidate implements Candidate interface. var _ Candidate = (*grokStreamCandidate)(nil) // String provides a string representation of the candidate. func (c *grokStreamCandidate) String() string { return fmt.Sprintf("StreamingCandidate(Index: %d, FinishReason: %s)", c.streamChoice.Index, c.streamChoice.FinishReason) } // Parts returns the parts of this streaming chunk candidate. func (c *grokStreamCandidate) Parts() []Part { var parts []Part // Include text content if present if c.streamChoice.Delta.Content != "" { parts = append(parts, &grokStreamPart{ content: c.streamChoice.Delta.Content, }) } // Include tool calls if present if len(c.streamChoice.Delta.ToolCalls) > 0 { // Convert ChatCompletionToolCallDelta to ChatCompletionMessageToolCall toolCalls := make([]openai.ChatCompletionMessageToolCall, 0, len(c.streamChoice.Delta.ToolCalls)) for _, delta := range c.streamChoice.Delta.ToolCalls { // Create a new ChatCompletionMessageToolCall directly toolCall := openai.ChatCompletionMessageToolCall{ ID: delta.ID, Function: openai.ChatCompletionMessageToolCallFunction{ Name: delta.Function.Name, Arguments: delta.Function.Arguments, }, Type: "function", // The type is always "function" for function calls } toolCalls = append(toolCalls, toolCall) } parts = append(parts, &grokStreamPart{ toolCalls: toolCalls, }) } return parts } // grokStreamPart adapts streaming parts to the Part interface. type grokStreamPart struct { content string toolCalls []openai.ChatCompletionMessageToolCall } // Ensure the streaming part implements Part interface. var _ Part = (*grokStreamPart)(nil) // AsText returns the text content of this part if it has any. func (p *grokStreamPart) AsText() (string, bool) { return p.content, p.content != "" } // AsFunctionCalls returns the function calls from this part if it has any. func (p *grokStreamPart) AsFunctionCalls() ([]FunctionCall, bool) { if len(p.toolCalls) == 0 { return nil, false } // Count valid function calls first validCount := 0 for _, tc := range p.toolCalls { // Only count tool calls that have a function name if tc.Function.Name != "" { validCount++ } } // If no valid function calls, return nil if validCount == 0 { return nil, false } // Create properly sized array completeCalls := make([]FunctionCall, 0, validCount) // Process tool calls for _, tc := range p.toolCalls { // Skip tool calls that don't have a complete function definition yet if tc.Function.Name == "" { continue } var args map[string]any // Attempt to unmarshal arguments if present if tc.Function.Arguments != "" { if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil { klog.V(2).Infof("Error unmarshaling function arguments: %v", err) // Continue with empty args if unmarshal fails args = make(map[string]any) } } else { // Initialize empty args map if no arguments provided args = make(map[string]any) } completeCalls = append(completeCalls, FunctionCall{ ID: tc.ID, Name: tc.Function.Name, Arguments: args, }) } return completeCalls, len(completeCalls) > 0 } ================================================ FILE: gollm/http_journal.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "bytes" "io" "net/http" "net/http/httputil" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal" "k8s.io/klog/v2" ) // journalingRoundTripper wraps an existing http.RoundTripper to record requests and responses. type journalingRoundTripper struct { next http.RoundTripper // The actual transport that does the network call } // RoundTrip satisfies the http.RoundTripper interface. It intercepts an HTTP request, // logs it, passes it to the next handler, and then logs the response. // It includes special handling to correctly parse and summarize streaming responses. func (jrt *journalingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { recorder := journal.RecorderFromContext(req.Context()) // Log the outgoing request. reqBytes, err := httputil.DumpRequestOut(req, true) if err == nil { err = recorder.Write(req.Context(), &journal.Event{ Action: journal.ActionHTTPRequest, Payload: map[string]any{"request": string(reqBytes)}, }) if err != nil { klog.Errorf("Error writing outgoing request to journal: %v", err) } } // Pass the request to the next RoundTripper to make the actual network call. resp, err := jrt.next.RoundTrip(req) if err != nil { writeErr := recorder.Write(req.Context(), &journal.Event{ Action: journal.ActionHTTPError, Payload: map[string]any{"error": "http transport failed", "detail": err.Error()}, }) if writeErr != nil { klog.Errorf("Error writing RoundTripper error to journal: %v", writeErr) } klog.Errorf("RoundTripper error: %v", err) return nil, err } // Read the entire response body so we can log it and then pass it along. bodyBytes, err := io.ReadAll(resp.Body) if err != nil { // handle error klog.Errorf("Error reading response body (for logging): %v", err) return nil, err } resp.Body.Close() // Close the original body // Default payload is the raw body, for non-streaming responses. logPayload := map[string]any{ "status": resp.Status, "headers": resp.Header, "body": string(bodyBytes), } // Write the final event to the journal. err = recorder.Write(req.Context(), &journal.Event{ Action: journal.ActionHTTPResponse, Payload: logPayload, }) if err != nil { // Log the error and continue klog.Errorf("Error writing to journal: %v", err) } // IMPORTANT: Return the original, untouched body to the client. resp.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) return resp, nil } // withJournaling is a decorator function that wraps an http.Client's transport // with the journalingRoundTripper, but only if a recorder is found in the context. func withJournaling(client *http.Client) *http.Client { // wrap the transport client.Transport = &journalingRoundTripper{ next: client.Transport, } return client } ================================================ FILE: gollm/interfaces.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "fmt" "io" "iter" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) // Client is a client for a language model. type Client interface { io.Closer // StartChat starts a new multi-turn chat with a language model. StartChat(systemPrompt, model string) Chat // GenerateCompletion generates a single completion for a given prompt. GenerateCompletion(ctx context.Context, req *CompletionRequest) (CompletionResponse, error) // SetResponseSchema constrains LLM responses to match the provided schema. // Calling with nil will clear the current schema. SetResponseSchema(schema *Schema) error // ListModels lists the models available in the LLM. ListModels(ctx context.Context) ([]string, error) } // Chat is an active conversation with a language model. // Messages are sent and received, and add to a conversation history. type Chat interface { // Send adds a user message to the chat, and gets the response from the LLM. // Note that this method automatically updates the state of the Chat, // you do not need to "replay" any messages from the LLM. Send(ctx context.Context, contents ...any) (ChatResponse, error) // SendStreaming is the streaming version of Send. SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) // SetFunctionDefinitions configures the set of tools (functions) available to the LLM // for function calling. SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error // IsRetryableError returns true if the error is retryable. IsRetryableError(error) bool // Initialize initializes the chat with a previous conversation history. Initialize(messages []*api.Message) error } // CompletionRequest is a request to generate a completion for a given prompt. type CompletionRequest struct { Model string `json:"model,omitempty"` Prompt string `json:"prompt,omitempty"` } // CompletionResponse is a response from the GenerateCompletion method. type CompletionResponse interface { Response() string UsageMetadata() any } // FunctionCall is a function call to a language model. // The LLM will reply with a FunctionCall to a user-defined function, and we will send the results back. type FunctionCall struct { ID string `json:"id,omitempty"` Name string `json:"name,omitempty"` Arguments map[string]any `json:"arguments,omitempty"` } // FunctionDefinition is a user-defined function that can be called by the LLM. // If the LLM determines the function should be called, it will reply with a FunctionCall object; // we will invoke the function and the results back. type FunctionDefinition struct { Name string `json:"name,omitempty"` Description string `json:"description,omitempty"` Parameters *Schema `json:"parameters,omitempty"` } // Schema is a schema for a function definition. type Schema struct { Type SchemaType `json:"type,omitempty"` Properties map[string]*Schema `json:"properties,omitempty"` Items *Schema `json:"items,omitempty"` Description string `json:"description,omitempty"` Required []string `json:"required,omitempty"` } // ToRawSchema converts a Schema to a json.RawMessage. func (s *Schema) ToRawSchema() (json.RawMessage, error) { jsonSchema, err := json.Marshal(s) if err != nil { return nil, fmt.Errorf("converting tool schema to json: %w", err) } var rawSchema json.RawMessage if err := json.Unmarshal(jsonSchema, &rawSchema); err != nil { return nil, fmt.Errorf("converting tool schema to json.RawMessage: %w", err) } return rawSchema, nil } // SchemaType is the type of a field in a Schema. type SchemaType string const ( TypeObject SchemaType = "object" TypeArray SchemaType = "array" TypeString SchemaType = "string" TypeBoolean SchemaType = "boolean" TypeNumber SchemaType = "number" TypeInteger SchemaType = "integer" ) // FunctionCallResult is the result of a function call. // We use this to send the results back to the LLM. type FunctionCallResult struct { ID string `json:"id,omitempty"` Name string `json:"name,omitempty"` Result map[string]any `json:"result,omitempty"` } // ChatResponse is a generic chat response from the LLM. type ChatResponse interface { UsageMetadata() any // Candidates are a set of candidate responses from the LLM. // The LLM may return multiple candidates, and we can choose the best one. Candidates() []Candidate } // ChatResponseIterator is a streaming chat response from the LLM. type ChatResponseIterator iter.Seq2[ChatResponse, error] // Candidate is one of a set of candidate response from the LLM. type Candidate interface { // String returns a string representation of the candidate. fmt.Stringer // Parts returns the parts of the candidate. Parts() []Part } // Part is a part of a candidate response from the LLM. // It can be a text response, or a function call. // A response may comprise multiple parts, // for example a text response and a function call // where the text response is "I need to do the necessary" // and then the function call is "do_necessary". type Part interface { // AsText returns the text of the part. // if the part is not text, it returns ("", false) AsText() (string, bool) // AsFunctionCalls returns the function calls of the part. // if the part is not a function call, it returns (nil, false) AsFunctionCalls() ([]FunctionCall, bool) } ================================================ FILE: gollm/llamacpp.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/url" "os" "k8s.io/klog/v2" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) func init() { if err := RegisterProvider("llamacpp", llamacppFactory); err != nil { klog.Fatalf("Failed to register llamacpp provider: %v", err) } } // llamacppFactory is the provider factory function for llama.cpp. // Supports ClientOptions for custom configuration, including skipVerifySSL. func llamacppFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewLlamaCppClient(ctx, opts) } type LlamaCppClient struct { baseURL *url.URL httpClient *http.Client responseSchema *llamacppSchema } type LlamaCppChat struct { client *LlamaCppClient model string history []llamacppChatMessage tools []llamacppTool } var _ Client = &LlamaCppClient{} // NewLlamaCppClient creates a new client for llama.cpp. // Supports custom HTTP client and skipVerifySSL via ClientOptions. func NewLlamaCppClient(ctx context.Context, opts ClientOptions) (*LlamaCppClient, error) { host := os.Getenv("LLAMACPP_HOST") if host == "" { host = "http://127.0.0.1:8080/" } baseURL, err := url.Parse(host) if err != nil { return nil, fmt.Errorf("parsing host %q: %w", host, err) } klog.Infof("using llama.cpp with base url %v", baseURL.String()) httpClient := createCustomHTTPClient(opts.SkipVerifySSL) return &LlamaCppClient{ baseURL: baseURL, httpClient: httpClient, }, nil } func (c *LlamaCppClient) Close() error { return nil } func (c *LlamaCppClient) GenerateCompletion(ctx context.Context, request *CompletionRequest) (CompletionResponse, error) { llamacppRequest := &llamacppCompletionRequest{ Prompt: request.Prompt, JSONSchema: c.responseSchema, } llamacppResponse, err := c.doCompletion(ctx, llamacppRequest) if err != nil { return nil, err } if llamacppResponse.Content == "" { return nil, fmt.Errorf("no response returned from llamacpp") } response := &LlamaCppCompletionResponse{llamacppResponse: llamacppResponse} return response, nil } func (c *LlamaCppClient) doRequest(ctx context.Context, httpMethod, relativePath string, req any, response any) error { body, err := json.Marshal(req) if err != nil { return fmt.Errorf("building json body: %w", err) } u := c.baseURL.JoinPath(relativePath) klog.V(2).Infof("sending %s request to %v: %v", httpMethod, u.String(), string(body)) httpRequest, err := http.NewRequestWithContext(ctx, httpMethod, u.String(), bytes.NewReader(body)) if err != nil { return fmt.Errorf("building http request: %w", err) } httpRequest.Header.Set("Content-Type", "application/json") httpResponse, err := c.httpClient.Do(httpRequest) if err != nil { return fmt.Errorf("performing http request: %w", err) } defer httpResponse.Body.Close() b, err := io.ReadAll(httpResponse.Body) if err != nil { return fmt.Errorf("reading response body: %w", err) } if httpResponse.StatusCode != 200 { return fmt.Errorf("unexpected http status: %q with response %q", httpResponse.Status, string(b)) } if err := json.Unmarshal(b, response); err != nil { return fmt.Errorf("unmarshalling json response: %w", err) } return nil } func (c *LlamaCppClient) doCompletion(ctx context.Context, req *llamacppCompletionRequest) (*llamacppCompletionResponse, error) { completionResponse := &llamacppCompletionResponse{} if err := c.doRequest(ctx, "POST", "completion", req, completionResponse); err != nil { return nil, err } return completionResponse, nil } func (c *LlamaCppClient) doChat(ctx context.Context, req *llamacppChatRequest) (*llamacppChatResponse, error) { chatResponse := &llamacppChatResponse{} if err := c.doRequest(ctx, "POST", "v1/chat/completions", req, chatResponse); err != nil { return nil, err } return chatResponse, nil } func (c *LlamaCppClient) ListModels(ctx context.Context) ([]string, error) { return nil, fmt.Errorf("model switching not supported by llama.cpp") } func (c *LlamaCppClient) SetResponseSchema(responseSchema *Schema) error { llamaSchema := toLlamacppSchema(responseSchema) c.responseSchema = llamaSchema return nil } func (c *LlamaCppClient) StartChat(systemPrompt, model string) Chat { return &LlamaCppChat{ client: c, model: model, history: []llamacppChatMessage{ { Role: "system", Content: ptrTo(systemPrompt), }, }, } } type LlamaCppCompletionResponse struct { llamacppResponse *llamacppCompletionResponse } func (r *LlamaCppCompletionResponse) Response() string { return r.llamacppResponse.Content } func (r *LlamaCppCompletionResponse) UsageMetadata() any { return nil } func (c *LlamaCppChat) Send(ctx context.Context, contents ...any) (ChatResponse, error) { log := klog.FromContext(ctx) for _, content := range contents { switch v := content.(type) { case string: message := llamacppChatMessage{ Role: "user", Content: ptrTo(v), } c.history = append(c.history, message) case FunctionCallResult: resultJSON, err := json.Marshal(v.Result) if err != nil { return nil, fmt.Errorf("marshalling function call result: %w", err) } message := llamacppChatMessage{ Role: "tool", // TODO: Do we need ToolCallID? ToolCallID: toolCallId, Content: ptrTo(string(resultJSON)), } c.history = append(c.history, message) default: return nil, fmt.Errorf("unsupported content type: %T", v) } } req := &llamacppChatRequest{ Model: c.model, Messages: c.history, // Stream: ptrTo(false), Tools: c.tools, } var llmacppResponse *LlamaCppChatResponse resp, err := c.client.doChat(ctx, req) if err != nil { return nil, err } log.V(2).Info("received response from llama.cpp", "resp", resp) llmacppResponse = &LlamaCppChatResponse{ LlamaCppResponse: *resp, } for i, choice := range resp.Choices { candidate := &LlamaCppCandidate{} if choice.Message != nil && choice.Message.Content != nil { parts := &LlamaCppPart{ text: *choice.Message.Content, } candidate.parts = append(candidate.parts, parts) } if choice.Message != nil && len(choice.Message.ToolCalls) != 0 { var functionCalls []FunctionCall for _, toolCall := range choice.Message.ToolCalls { functionCall := FunctionCall{ Name: toolCall.Function.Name, } if toolCall.Function.Arguments != "" { arguments := make(map[string]any) if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil { return nil, fmt.Errorf("parsing function call arguments: %w", err) } functionCall.Arguments = arguments } functionCalls = append(functionCalls, functionCall) } parts := &LlamaCppPart{ functionCalls: functionCalls, } candidate.parts = append(candidate.parts, parts) } llmacppResponse.candidates = append(llmacppResponse.candidates, candidate) if i == 0 { if choice.Message != nil { msg := llamacppChatMessage{ Role: "assistant", Content: choice.Message.Content, ToolCalls: choice.Message.ToolCalls, ToolCallID: choice.Message.ToolCallID, } c.history = append(c.history, msg) } } } return llmacppResponse, nil } func (c *LlamaCppChat) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { // TODO: Implement streaming response, err := c.Send(ctx, contents...) if err != nil { return nil, err } return singletonChatResponseIterator(response), nil } func (c *LlamaCppChat) IsRetryableError(err error) bool { // TODO(droot): Implement this return false } func (c *LlamaCppChat) Initialize(messages []*api.Message) error { klog.Warning("chat history persistence is not supported for provider 'llamacpp', using in-memory chat history") return nil } func ptrTo[T any](t T) *T { return &t } type LlamaCppChatResponse struct { candidates []*LlamaCppCandidate LlamaCppResponse llamacppChatResponse } var _ ChatResponse = &LlamaCppChatResponse{} func (r *LlamaCppChatResponse) MarshalJSON() ([]byte, error) { formatted := RecordChatResponse{ Raw: r.LlamaCppResponse, } return json.Marshal(&formatted) } func (r *LlamaCppChatResponse) String() string { return fmt.Sprintf("LlamaCppChatResponse{candidates=%v}", r.candidates) } // func (r *LlamaCppChatResponse) String() string { // var sb strings.Builder // fmt.Fprintf(&sb, "LlamaCppChatResponse{candidates=[") // for _, candidate := range r.candidates { // fmt.Fprintf(&sb, "%v", candidate) // } // fmt.Fprintf(&sb, "]}") // return sb.String() // } func (r *LlamaCppChatResponse) UsageMetadata() any { return nil } func (r *LlamaCppChatResponse) Candidates() []Candidate { var cads []Candidate for _, candidate := range r.candidates { cads = append(cads, candidate) } return cads } type LlamaCppCandidate struct { parts []*LlamaCppPart } func (r *LlamaCppCandidate) String() string { return r.parts[0].text } func (r *LlamaCppCandidate) Parts() []Part { var out []Part for _, part := range r.parts { out = append(out, part) } return out } type LlamaCppPart struct { text string functionCalls []FunctionCall } func (p *LlamaCppPart) AsText() (string, bool) { if len(p.text) > 0 { return p.text, true } return "", false } func (p *LlamaCppPart) AsFunctionCalls() ([]FunctionCall, bool) { if len(p.functionCalls) > 0 { return p.functionCalls, true } return nil, false } func (c *LlamaCppChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error { var tools []llamacppTool for _, functionDefinition := range functionDefinitions { tools = append(tools, toLlamacppTool(functionDefinition)) } c.tools = tools return nil } func toLlamacppTool(fnDef *FunctionDefinition) llamacppTool { function := &llamacppFunction{ Description: fnDef.Description, Name: fnDef.Name, } if fnDef.Parameters != nil { function.Parameters = toLlamacppSchema(fnDef.Parameters) } tool := llamacppTool{ Type: "function", Function: function, } return tool } func toLlamacppSchema(in *Schema) *llamacppSchema { if in == nil { return nil } out := &llamacppSchema{ Type: string(in.Type), Items: toLlamacppSchema(in.Items), Description: in.Description, Required: in.Required, } if in.Properties != nil { out.Properties = make(map[string]llamacppSchema, len(in.Properties)) for k, v := range in.Properties { out.Properties[k] = *toLlamacppSchema(v) } } return out } type llamacppCompletionRequest struct { // See https://github.com/ggerganov/llama.cpp/blob/master/examples/server/README.md#post-completion-given-a-prompt-it-returns-the-predicted-completion Prompt string `json:"prompt,omitempty"` JSONSchema *llamacppSchema `json:"json_schema,omitempty"` } type llamacppCompletionResponse struct { Content string `json:"content,omitempty"` Index int32 `json:"index,omitempty"` IDSlot int32 `json:"id_slot,omitempty"` Stop bool `json:"stop,omitempty"` Model string `json:"model,omitempty"` TokensPredicted int32 `json:"tokens_predicted,omitempty"` TokensEvaluated int32 `json:"tokens_evaluated,omitempty"` // "generation_settings":{"n_predict":-1,"seed":4294967295,"temperature":0.800000011920929,"dynatemp_range":0.0,"dynatemp_exponent":1.0,"top_k":40,"top_p":0.949999988079071,"min_p":0.05000000074505806,"xtc_probability":0.0,"xtc_threshold":0.10000000149011612,"typical_p":1.0,"repeat_last_n":64,"repeat_penalty":1.0,"presence_penalty":0.0,"frequency_penalty":0.0,"dry_multiplier":0.0,"dry_base":1.75,"dry_allowed_length":2,"dry_penalty_last_n":16384,"dry_sequence_breakers":["\n",":","\"","*"],"mirostat":0,"mirostat_tau":5.0,"mirostat_eta":0.10000000149011612,"stop":[],"max_tokens":-1,"n_keep":0,"n_discard":0,"ignore_eos":false,"stream":false,"logit_bias":[],"n_probs":0,"min_keep":0,"grammar":"","grammar_lazy":false,"grammar_triggers":[],"preserved_tokens":[],"chat_format":"Content-only","samplers":["penalties","dry","top_k","typ_p","top_p","min_p","xtc","temperature"],"speculative.n_max":16,"speculative.n_min":0,"speculative.p_min":0.75,"timings_per_token":false,"post_sampling_probs":false,"lora":[]}, // GenerationSettings llamacppGenerationSettings `json:"generation_settings,omitempty"` Prompt string `json:"prompt,omitempty"` HasNewLine bool `json:"has_new_line,omitempty"` Truncated bool `json:"truncated,omitempty"` StopType string `json:"stop_type,omitempty"` StoppingWord string `json:"stopping_word,omitempty"` TokensCached int32 `json:"tokens_cached,omitempty"` Timings llamacppTimings `json:"timings,omitempty"` } type llamacppTimings struct { PromptN int32 `json:"prompt_n,omitempty"` PromptMs float64 `json:"prompt_ms,omitempty"` PromptPerTokenMs float64 `json:"prompt_per_token_ms,omitempty"` PromptPerSecond float64 `json:"prompt_per_second,omitempty"` PredictedN int32 `json:"predicted_n,omitempty"` PredictedMs float64 `json:"predicted_ms,omitempty"` PredictedPerTokenMs float64 `json:"predicted_per_token_ms,omitempty"` PredictedPerSecond float64 `json:"predicted_per_second,omitempty"` } type llamacppChatRequest struct { Model string `json:"model,omitempty"` Messages []llamacppChatMessage `json:"messages,omitempty"` Tools []llamacppTool `json:"tools,omitempty"` } type llamacppChatResponse struct { Choices []llamacppChoice `json:"choices,omitempty"` Created int64 `json:"created,omitempty"` Model string `json:"model,omitempty"` SystemFingerprint string `json:"system_fingerprint,omitempty"` Object string `json:"object,omitempty"` Usage *llamacppUsage `json:"usage,omitempty"` Id string `json:"id,omitempty"` Timings *llamacppTimings `json:"timings,omitempty"` } type llamacppChoice struct { FinishReason string `json:"finish_reason,omitempty"` Index int32 `json:"index,omitempty"` Message *llamacppChatMessage `json:"message,omitempty"` } type llamacppUsage struct { CompletionTokens int32 `json:"completion_tokens,omitempty"` PromptTokens int32 `json:"prompt_tokens,omitempty"` TotalTokens int32 `json:"total_tokens,omitempty"` } type llamacppChatMessage struct { Role string `json:"role,omitempty"` Content *string `json:"content,omitempty"` ToolCalls []llamacppToolCall `json:"tool_calls,omitempty"` ToolCallID string `json:"tool_call_id,omitempty"` } type llamacppToolCall struct { Type string `json:"type,omitempty"` Function llamacppFunctionCall `json:"function,omitempty"` } type llamacppFunctionCall struct { Name string `json:"name,omitempty"` Arguments string `json:"arguments,omitempty"` ID string `json:"id,omitempty"` } type llamacppTool struct { Type string `json:"type,omitempty"` Function *llamacppFunction `json:"function,omitempty"` } type llamacppFunction struct { Description string `json:"description,omitempty"` Name string `json:"name,omitempty"` Parameters *llamacppSchema `json:"parameters,omitempty"` } type llamacppSchema struct { Type string `json:"type,omitempty"` Required []string `json:"required,omitempty"` Items *llamacppSchema `json:"items,omitempty"` Properties map[string]llamacppSchema `json:"properties,omitempty"` Description string `json:"description,omitempty"` Enum []string `json:"enum,omitempty"` } ================================================ FILE: gollm/ollama.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "fmt" "github.com/ollama/ollama/api" "github.com/ollama/ollama/envconfig" "k8s.io/klog/v2" kctlApi "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) func init() { if err := RegisterProvider("ollama", ollamaFactory); err != nil { klog.Fatalf("Failed to register ollama provider: %v", err) } } // ollamaFactory is the provider factory function for Ollama. // Supports ClientOptions for custom configuration, including skipVerifySSL. func ollamaFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewOllamaClient(ctx, opts) } const ( defaultOllamaModel = "gemma3:latest" ) type OllamaClient struct { client *api.Client } type OllamaChat struct { client *api.Client model string history []api.Message tools []api.Tool } var _ Client = &OllamaClient{} // NewOllamaClient creates a new client for Ollama. // Supports custom HTTP client and skipVerifySSL via ClientOptions if the SDK supports it. func NewOllamaClient(ctx context.Context, opts ClientOptions) (*OllamaClient, error) { // Create custom HTTP client with SSL verification option from client options httpClient := createCustomHTTPClient(opts.SkipVerifySSL) client := api.NewClient(envconfig.Host(), httpClient) return &OllamaClient{ client: client, }, nil } func (c *OllamaClient) Close() error { return nil } func (c *OllamaClient) GenerateCompletion(ctx context.Context, request *CompletionRequest) (CompletionResponse, error) { req := &api.GenerateRequest{ Model: request.Model, Prompt: request.Prompt, Stream: ptrTo(false), } var ollamaResponse *OllamaCompletionResponse respFunc := func(resp api.GenerateResponse) error { ollamaResponse = &OllamaCompletionResponse{response: resp.Response} return nil } err := c.client.Generate(ctx, req, respFunc) if err != nil { return nil, err } return ollamaResponse, nil } func (c *OllamaClient) ListModels(ctx context.Context) ([]string, error) { modelResponse, err := c.client.List(ctx) if err != nil { return nil, err } var models []string for _, model := range modelResponse.Models { models = append(models, model.Name) } return models, nil } func (c *OllamaClient) SetResponseSchema(schema *Schema) error { return nil } func (c *OllamaClient) StartChat(systemPrompt, model string) Chat { return &OllamaChat{ client: c.client, model: model, history: []api.Message{ { Role: "system", Content: systemPrompt, }, }, } } type OllamaCompletionResponse struct { response string } func (r *OllamaCompletionResponse) Response() string { return r.response } func (r *OllamaCompletionResponse) UsageMetadata() any { return nil } func (c *OllamaChat) Send(ctx context.Context, contents ...any) (ChatResponse, error) { log := klog.FromContext(ctx) for _, content := range contents { switch v := content.(type) { case string: message := api.Message{ Role: "user", Content: v, } c.history = append(c.history, message) case FunctionCallResult: message := api.Message{ Role: "user", Content: fmt.Sprintf("Function call result: %s", v.Result), } c.history = append(c.history, message) default: return nil, fmt.Errorf("unsupported content type: %T", v) } } req := &api.ChatRequest{ Model: c.model, Messages: c.history, // set streaming to false Stream: new(bool), Tools: c.tools, } var ollamaResponse *OllamaChatResponse respFunc := func(resp api.ChatResponse) error { log.Info("received response from ollama", "resp", resp) ollamaResponse = &OllamaChatResponse{ ollamaResponse: resp, candidates: []*OllamaCandidate{ { parts: []OllamaPart{ { text: resp.Message.Content, toolCalls: resp.Message.ToolCalls, }, }, }, }, } c.history = append(c.history, resp.Message) return nil } err := c.client.Chat(ctx, req, respFunc) if err != nil { return nil, err } log.Info("ollama response", "parsed_response", ollamaResponse) return ollamaResponse, nil } func (c *OllamaChat) IsRetryableError(err error) bool { // TODO(droot): Implement this return false } func (c *OllamaChat) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { // TODO: Implement streaming response, err := c.Send(ctx, contents...) if err != nil { return nil, err } return singletonChatResponseIterator(response), nil } func (c *OllamaChat) Initialize(messages []*kctlApi.Message) error { klog.Warning("chat history persistence is not supported for provider 'ollama', using in-memory chat history") return nil } type OllamaChatResponse struct { candidates []*OllamaCandidate ollamaResponse api.ChatResponse } var _ ChatResponse = &OllamaChatResponse{} func (r *OllamaChatResponse) MarshalJSON() ([]byte, error) { formatted := RecordChatResponse{ Raw: r.ollamaResponse, } return json.Marshal(&formatted) } func (r *OllamaChatResponse) String() string { return fmt.Sprintf("OllamaChatResponse{candidates=%v}", r.candidates) } func (r *OllamaChatResponse) UsageMetadata() any { return nil } func (r *OllamaChatResponse) Candidates() []Candidate { var cads []Candidate for _, candidate := range r.candidates { cads = append(cads, candidate) } return cads } type OllamaCandidate struct { parts []OllamaPart } func (r *OllamaCandidate) String() string { return r.parts[0].text } func (r *OllamaCandidate) Parts() []Part { var parts []Part for _, part := range r.parts { parts = append(parts, &OllamaPart{ text: part.text, toolCalls: part.toolCalls, }) } return parts } type OllamaPart struct { text string toolCalls []api.ToolCall } func (p *OllamaPart) AsText() (string, bool) { if len(p.text) > 0 { return p.text, true } return "", false } func (p *OllamaPart) AsFunctionCalls() ([]FunctionCall, bool) { if len(p.toolCalls) > 0 { var functionCalls []FunctionCall for _, toolCall := range p.toolCalls { functionCalls = append(functionCalls, FunctionCall{ Name: toolCall.Function.Name, Arguments: toolCall.Function.Arguments, }) } return functionCalls, true } return nil, false } func (c *OllamaChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error { var tools []api.Tool for _, functionDefinition := range functionDefinitions { tools = append(tools, fnDefToOllamaTool(functionDefinition)) } c.tools = tools return nil } func fnDefToOllamaTool(fnDef *FunctionDefinition) api.Tool { tool := api.Tool{ Type: "function", Function: api.ToolFunction{ Name: fnDef.Name, Description: fnDef.Description, Parameters: struct { Type string `json:"type"` Required []string `json:"required"` Properties map[string]struct { Type string `json:"type"` Description string `json:"description"` Enum []string `json:"enum,omitempty"` } `json:"properties"` }{ Type: "object", Required: fnDef.Parameters.Required, Properties: map[string]struct { Type string `json:"type"` Description string `json:"description"` Enum []string `json:"enum,omitempty"` }{}, }, }, } for paramName, param := range fnDef.Parameters.Properties { tool.Function.Parameters.Properties[paramName] = struct { Type string `json:"type"` Description string `json:"description"` Enum []string `json:"enum,omitempty"` }{ Type: string(param.Type), Description: param.Description, } } return tool } ================================================ FILE: gollm/openai.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "errors" "fmt" "os" "strings" openai "github.com/openai/openai-go" "github.com/openai/openai-go/option" "github.com/openai/openai-go/responses" "k8s.io/klog/v2" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) // Package-level env var storage (OpenAI env) var ( openAIAPIKey string openAIEndpoint string openAIAPIBase string openAIModel string openAIUseResponsesAPI bool ) // init reads and caches OpenAI environment variables: // - OPENAI_API_KEY, OPENAI_ENDPOINT, OPENAI_API_BASE, OPENAI_MODEL // // These serve as defaults; the model can be overridden by the Cobra --model flag. // After loading env values, it registers the OpenAI provider factory. func init() { // Load environment variables openAIAPIKey = os.Getenv("OPENAI_API_KEY") openAIEndpoint = os.Getenv("OPENAI_ENDPOINT") openAIAPIBase = os.Getenv("OPENAI_API_BASE") openAIModel = os.Getenv("OPENAI_MODEL") if val := os.Getenv("OPENAI_USE_RESPONSES_API"); strings.ToLower(val) == "true" { openAIUseResponsesAPI = true klog.InfoS("Using responses API for openai", "baseURL", openAIAPIBase, "endpoint", openAIEndpoint, "model", openAIModel) } // Register "openai" as the provider ID if err := RegisterProvider("openai", newOpenAIClientFactory); err != nil { klog.Fatalf("Failed to register openai provider: %v", err) } // Also register with any aliases defined in config aliases := []string{"openai-compatible"} for _, alias := range aliases { if err := RegisterProvider(alias, newOpenAIClientFactory); err != nil { klog.Warningf("Failed to register openai provider alias %q: %v", alias, err) } } } // OpenAIClient implements the gollm.Client interface for OpenAI models. type OpenAIClient struct { client openai.Client } // Ensure OpenAIClient implements the Client interface. var _ Client = &OpenAIClient{} // NewOpenAIClient creates a new client for interacting with OpenAI. // Supports custom HTTP client (e.g., for skipping SSL verification). func NewOpenAIClient(ctx context.Context, opts ClientOptions) (*OpenAIClient, error) { // Get API key from loaded env var apiKey := openAIAPIKey if apiKey == "" { return nil, errors.New("OpenAI API key not found. Set via OPENAI_API_KEY env var") } // Set options for client creation options := []option.RequestOption{option.WithAPIKey(apiKey)} // Check for custom endpoint or API base URL baseURL := openAIEndpoint if baseURL == "" { baseURL = openAIAPIBase } if baseURL != "" { klog.Infof("Using custom OpenAI base URL: %s", baseURL) options = append(options, option.WithBaseURL(baseURL)) } // Support custom HTTP client (e.g., skip SSL verification) httpClient := createCustomHTTPClient(opts.SkipVerifySSL) httpClient = withJournaling(httpClient) options = append(options, option.WithHTTPClient(httpClient)) return &OpenAIClient{ client: openai.NewClient(options...), }, nil } // Close cleans up any resources used by the client. func (c *OpenAIClient) Close() error { // No specific cleanup needed for the OpenAI client currently. return nil } // StartChat starts a new chat session. func (c *OpenAIClient) StartChat(systemPrompt, model string) Chat { // Get the model to use for this chat selectedModel := getOpenAIModel(model) klog.V(1).Infof("Starting new OpenAI chat session with model: %s", selectedModel) if openAIUseResponsesAPI { // Initialize history with system prompt if provided history := responses.ResponseInputParam{} if systemPrompt != "" { history = append(history, responses.ResponseInputItemUnionParam{ OfMessage: &responses.EasyInputMessageParam{ Content: responses.EasyInputMessageContentUnionParam{ OfString: openai.String(systemPrompt), }, Role: responses.EasyInputMessageRoleSystem, }, }) } return &openAIResponseChatSession{ client: c.client, history: history, model: selectedModel, // functionDefinitions and tools will be set later via SetFunctionDefinitions params: responses.ResponseNewParams{ Model: selectedModel, Temperature: openai.Float(0.2), MaxOutputTokens: openai.Int(2048), Reasoning: responses.ReasoningParam{ Effort: responses.ReasoningEffortLow, }, Store: openai.Bool(false), }, } } // by default use completion endpoint // Initialize history with system prompt if provided history := []openai.ChatCompletionMessageParamUnion{} if systemPrompt != "" { history = append(history, openai.SystemMessage(systemPrompt)) } return &openAIChatSession{ client: c.client, history: history, model: selectedModel, // functionDefinitions and tools will be set later via SetFunctionDefinitions } } // simpleCompletionResponse is a basic implementation of CompletionResponse. type simpleCompletionResponse struct { content string } // Response returns the completion content. func (r *simpleCompletionResponse) Response() string { return r.content } // UsageMetadata returns nil for now. func (r *simpleCompletionResponse) UsageMetadata() any { return nil } // GenerateCompletion sends a completion request to the OpenAI API. func (c *OpenAIClient) GenerateCompletion(ctx context.Context, req *CompletionRequest) (CompletionResponse, error) { klog.Infof("OpenAI GenerateCompletion called with model: %s", req.Model) klog.V(1).Infof("Prompt:\n%s", req.Prompt) // Use the Chat Completions API with the new v1.0.0 API completion, err := c.client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{ Model: openai.ChatModel(req.Model), Messages: []openai.ChatCompletionMessageParamUnion{ openai.UserMessage(req.Prompt), }, }) if err != nil { return nil, fmt.Errorf("failed to generate OpenAI completion: %w", err) } // Check if there are choices and a message if len(completion.Choices) == 0 || completion.Choices[0].Message.Content == "" { return nil, errors.New("received an empty response from OpenAI") } // Return the content of the first choice resp := &simpleCompletionResponse{ content: completion.Choices[0].Message.Content, } return resp, nil } // SetResponseSchema is not implemented yet. func (c *OpenAIClient) SetResponseSchema(schema *Schema) error { klog.Warning("OpenAIClient.SetResponseSchema is not implemented yet") return nil } // ListModels returns a slice of strings with model IDs. // Note: This may not work with all OpenAI-compatible providers if they don't fully implement // the Models.List endpoint or return data in a different format. func (c *OpenAIClient) ListModels(ctx context.Context) ([]string, error) { res, err := c.client.Models.List(ctx) if err != nil { return nil, fmt.Errorf("error listing models from OpenAI: %w", err) } modelIDs := make([]string, 0, len(res.Data)) for _, model := range res.Data { modelIDs = append(modelIDs, model.ID) } return modelIDs, nil } // Chat Session Implementation type openAIChatSession struct { client openai.Client history []openai.ChatCompletionMessageParamUnion model string functionDefinitions []*FunctionDefinition // Stored in gollm format tools []openai.ChatCompletionToolParam // Stored in OpenAI format } // Ensure openAIChatSession implements the Chat interface. var _ Chat = (*openAIChatSession)(nil) // SetFunctionDefinitions stores the function definitions and converts them to OpenAI format. func (cs *openAIChatSession) SetFunctionDefinitions(defs []*FunctionDefinition) error { cs.functionDefinitions = defs cs.tools = nil // Clear previous tools if len(defs) > 0 { cs.tools = make([]openai.ChatCompletionToolParam, len(defs)) for i, gollmDef := range defs { klog.Infof("Processing function definition: %s", gollmDef.Name) // Process function parameters params, err := cs.convertFunctionParameters(gollmDef) if err != nil { return fmt.Errorf("failed to process parameters for function %s: %w", gollmDef.Name, err) } cs.tools[i] = openai.ChatCompletionToolParam{ Function: openai.FunctionDefinitionParam{ Name: gollmDef.Name, Description: openai.String(gollmDef.Description), Parameters: params, }, } } } klog.V(1).Infof("Set %d function definitions for OpenAI chat session", len(cs.functionDefinitions)) return nil } // Send sends the user message(s), appends to history, and gets the LLM response. func (cs *openAIChatSession) Send(ctx context.Context, contents ...any) (ChatResponse, error) { klog.V(1).InfoS("openAIChatSession.Send called", "model", cs.model, "history_len", len(cs.history)) // Process and append messages to history if err := cs.addContentsToHistory(contents); err != nil { return nil, err } // Prepare and send API request chatReq := openai.ChatCompletionNewParams{ Model: openai.ChatModel(cs.model), Messages: cs.history, } if len(cs.tools) > 0 { chatReq.Tools = cs.tools } // Call the OpenAI API klog.V(1).InfoS("Sending request to OpenAI Chat API", "model", cs.model, "messages", len(chatReq.Messages), "tools", len(chatReq.Tools)) completion, err := cs.client.Chat.Completions.New(ctx, chatReq) if err != nil { // TODO: Check if error is retryable using cs.IsRetryableError klog.Errorf("OpenAI ChatCompletion API error: %v", err) return nil, fmt.Errorf("OpenAI chat completion failed: %w", err) } klog.V(1).InfoS("Received response from OpenAI Chat API", "id", completion.ID, "choices", len(completion.Choices)) // Process the response if len(completion.Choices) == 0 { klog.Warning("Received response with no choices from OpenAI") return nil, errors.New("received empty response from OpenAI (no choices)") } // Add assistant's response (first choice) to history assistantMsg := completion.Choices[0].Message // Convert to param type before appending to history cs.history = append(cs.history, assistantMsg.ToParam()) klog.V(2).InfoS("Added assistant message to history", "content_present", assistantMsg.Content != "", "tool_calls", len(assistantMsg.ToolCalls)) // Wrap the response resp := &openAIChatResponse{ openaiCompletion: completion, } return resp, nil } // SendStreaming sends the user message(s) and returns an iterator for the LLM response stream. func (cs *openAIChatSession) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { klog.V(1).InfoS("Starting OpenAI streaming request", "model", cs.model) // Process and append messages to history if err := cs.addContentsToHistory(contents); err != nil { return nil, err } // Prepare and send API request chatReq := openai.ChatCompletionNewParams{ Model: openai.ChatModel(cs.model), Messages: cs.history, } if len(cs.tools) > 0 { chatReq.Tools = cs.tools } // Start the OpenAI streaming request klog.V(1).InfoS("Sending streaming request to OpenAI API", "model", cs.model, "messageCount", len(chatReq.Messages), "toolCount", len(chatReq.Tools)) stream := cs.client.Chat.Completions.NewStreaming(ctx, chatReq) // Create an accumulator to track the full response acc := openai.ChatCompletionAccumulator{} // Create and return the stream iterator return func(yield func(ChatResponse, error) bool) { defer stream.Close() var lastResponseChunk *openAIChatStreamResponse var currentContent strings.Builder var currentToolCalls []openai.ChatCompletionMessageToolCall // Process stream chunks for stream.Next() { chunk := stream.Current() // Update the accumulator with the new chunk acc.AddChunk(chunk) // Handle content completion if _, ok := acc.JustFinishedContent(); ok { klog.V(2).Info("Content stream finished") } // Handle refusal completion if refusal, ok := acc.JustFinishedRefusal(); ok { klog.V(2).Infof("Refusal stream finished: %v", refusal) yield(nil, fmt.Errorf("model refused to respond: %v", refusal)) return } // Handle tool call completion var toolCallsForThisChunk []openai.ChatCompletionMessageToolCall if tool, ok := acc.JustFinishedToolCall(); ok { klog.V(2).Infof("Tool call finished: %s %s", tool.Name, tool.Arguments) newToolCall := openai.ChatCompletionMessageToolCall{ ID: tool.ID, Function: openai.ChatCompletionMessageToolCallFunction{ Name: tool.Name, Arguments: tool.Arguments, }, } currentToolCalls = append(currentToolCalls, newToolCall) // Only include the newly finished tool call in this chunk toolCallsForThisChunk = []openai.ChatCompletionMessageToolCall{newToolCall} } streamResponse := &openAIChatStreamResponse{ streamChunk: chunk, accumulator: acc, content: "", // Default to empty content toolCalls: toolCallsForThisChunk, } // Only process content if there are choices and a delta if len(chunk.Choices) > 0 { delta := chunk.Choices[0].Delta if delta.Content != "" { currentContent.WriteString(delta.Content) streamResponse.content = delta.Content // Only set content if there's new content } } // Keep track of the last response for history lastResponseChunk = &openAIChatStreamResponse{ streamChunk: chunk, accumulator: acc, content: currentContent.String(), // Full accumulated content for history toolCalls: currentToolCalls, } // Only yield if there's actual content or tool calls to report if streamResponse.content != "" || len(streamResponse.toolCalls) > 0 { if !yield(streamResponse, nil) { return } } } // Check for errors after streaming completes if err := stream.Err(); err != nil { klog.Errorf("Error in OpenAI streaming: %v", err) yield(nil, fmt.Errorf("OpenAI streaming error: %w", err)) return } // Update conversation history with the complete message if lastResponseChunk != nil { completeMessage := openai.ChatCompletionMessage{ Content: currentContent.String(), Role: "assistant", ToolCalls: currentToolCalls, } // Append the full assistant response to history cs.history = append(cs.history, completeMessage.ToParam()) klog.V(2).InfoS("Added complete assistant message to history", "content_present", completeMessage.Content != "", "tool_calls", len(completeMessage.ToolCalls)) } }, nil } // IsRetryableError determines if an error from the OpenAI API should be retried. func (cs *openAIChatSession) IsRetryableError(err error) bool { if err == nil { return false } return DefaultIsRetryableError(err) } func (cs *openAIChatSession) Initialize(messages []*api.Message) error { klog.Warning("chat history persistence is not supported for provider 'openai', using in-memory chat history") return nil } // Helper structs for ChatResponse interface type openAIChatResponse struct { openaiCompletion *openai.ChatCompletion } var _ ChatResponse = (*openAIChatResponse)(nil) func (r *openAIChatResponse) UsageMetadata() any { // Check if the main completion object and Usage exist if r.openaiCompletion != nil && r.openaiCompletion.Usage.TotalTokens > 0 { // Check a field within Usage return r.openaiCompletion.Usage } return nil } func (r *openAIChatResponse) Candidates() []Candidate { if r.openaiCompletion == nil { return nil } candidates := make([]Candidate, len(r.openaiCompletion.Choices)) for i, choice := range r.openaiCompletion.Choices { candidates[i] = &openAICandidate{openaiChoice: &choice} } return candidates } type openAICandidate struct { openaiChoice *openai.ChatCompletionChoice } var _ Candidate = (*openAICandidate)(nil) func (c *openAICandidate) Parts() []Part { // Check if the choice exists before accessing Message if c.openaiChoice == nil { return nil } // OpenAI message can have Content AND ToolCalls var parts []Part if c.openaiChoice.Message.Content != "" { parts = append(parts, &openAIPart{content: c.openaiChoice.Message.Content}) } if len(c.openaiChoice.Message.ToolCalls) > 0 { parts = append(parts, &openAIPart{toolCalls: c.openaiChoice.Message.ToolCalls}) } return parts } // String provides a simple string representation for logging/debugging. func (c *openAICandidate) String() string { if c.openaiChoice == nil { return "" } content := "" if c.openaiChoice.Message.Content != "" { content = c.openaiChoice.Message.Content } toolCalls := len(c.openaiChoice.Message.ToolCalls) finishReason := string(c.openaiChoice.FinishReason) return fmt.Sprintf("Candidate(FinishReason: %s, ToolCalls: %d, Content: %q)", finishReason, toolCalls, content) } type openAIPart struct { content string toolCalls []openai.ChatCompletionMessageToolCall // Correct type } var _ Part = (*openAIPart)(nil) func (p *openAIPart) AsText() (string, bool) { return p.content, p.content != "" } func (p *openAIPart) AsFunctionCalls() ([]FunctionCall, bool) { return convertToolCallsToFunctionCalls(p.toolCalls) } // Update openAIChatStreamResponse to include accumulated content type openAIChatStreamResponse struct { streamChunk openai.ChatCompletionChunk accumulator openai.ChatCompletionAccumulator content string toolCalls []openai.ChatCompletionMessageToolCall } // Update Candidates() to use accumulated content func (r *openAIChatStreamResponse) Candidates() []Candidate { if len(r.streamChunk.Choices) == 0 { return nil } candidates := make([]Candidate, len(r.streamChunk.Choices)) for i, choice := range r.streamChunk.Choices { candidates[i] = &openAIStreamCandidate{ streamChoice: choice, content: r.content, toolCalls: r.toolCalls, } } return candidates } // Update openAIStreamCandidate to handle delta content type openAIStreamCandidate struct { streamChoice openai.ChatCompletionChunkChoice content string // This will now be just the delta content toolCalls []openai.ChatCompletionMessageToolCall } // Update Parts() to handle delta content func (c *openAIStreamCandidate) Parts() []Part { var parts []Part // Only include the delta content if c.content != "" { parts = append(parts, &openAIStreamPart{ content: c.content, }) } // Include accumulated tool calls if len(c.toolCalls) > 0 { parts = append(parts, &openAIStreamPart{ toolCalls: c.toolCalls, }) } return parts } // Add UsageMetadata implementation func (r *openAIChatStreamResponse) UsageMetadata() any { if r.accumulator.Usage.TotalTokens > 0 { return r.accumulator.Usage } return nil } // Add String implementation func (c *openAIStreamCandidate) String() string { return fmt.Sprintf("StreamingCandidate(Content: %q, ToolCalls: %d)", c.content, len(c.toolCalls)) } // Define openAIStreamPart type openAIStreamPart struct { content string toolCalls []openai.ChatCompletionMessageToolCall } // Ensure openAIStreamPart implements Part interface var _ Part = (*openAIStreamPart)(nil) func (p *openAIStreamPart) AsText() (string, bool) { return p.content, p.content != "" } func (p *openAIStreamPart) AsFunctionCalls() ([]FunctionCall, bool) { return convertToolCallsToFunctionCalls(p.toolCalls) } // convertSchemaForOpenAI converts and transforms a schema for OpenAI compatibility // This function handles both gollm Schema objects and ensures the final JSON meets OpenAI requirements func convertSchemaForOpenAI(schema *Schema) (*Schema, error) { if schema == nil { // Return a minimal valid object schema for OpenAI return &Schema{ Type: TypeObject, Properties: make(map[string]*Schema), }, nil } // Create a deep copy to avoid modifying the original validated := &Schema{ Description: schema.Description, Required: make([]string, len(schema.Required)), } copy(validated.Required, schema.Required) // Handle type validation and normalization based on OpenAI requirements switch schema.Type { case TypeObject: validated.Type = TypeObject // Objects MUST have properties for OpenAI (even if empty) validated.Properties = make(map[string]*Schema) if schema.Properties != nil { for key, prop := range schema.Properties { validatedProp, err := convertSchemaForOpenAI(prop) if err != nil { return nil, fmt.Errorf("validating property %q: %w", key, err) } validated.Properties[key] = validatedProp } } case TypeArray: validated.Type = TypeArray // Arrays MUST have items schema for OpenAI if schema.Items != nil { validatedItems, err := convertSchemaForOpenAI(schema.Items) if err != nil { return nil, fmt.Errorf("validating array items: %w", err) } validated.Items = validatedItems } else { // Default to string items if not specified validated.Items = &Schema{Type: TypeString} } case TypeString: validated.Type = TypeString case TypeNumber: validated.Type = TypeNumber case TypeInteger: // OpenAI prefers "number" for integers validated.Type = TypeNumber case TypeBoolean: validated.Type = TypeBoolean case "": // If no type specified, default to object with empty properties klog.Warningf("Schema has no type, defaulting to object") validated.Type = TypeObject validated.Properties = make(map[string]*Schema) default: // For unknown types, log a warning and default to object klog.Warningf("Unknown schema type '%s', defaulting to object", schema.Type) validated.Type = TypeObject validated.Properties = make(map[string]*Schema) } // Final validation: Ensure object types always have properties // This handles edge cases where malformed schemas might slip through if validated.Type == TypeObject && validated.Properties == nil { klog.Warningf("Object schema missing properties, initializing empty properties map") validated.Properties = make(map[string]*Schema) } return validated, nil } // convertFunctionParameters handles the conversion of gollm parameters to OpenAI format func (cs *openAIChatSession) convertFunctionParameters(gollmDef *FunctionDefinition) (openai.FunctionParameters, error) { var params openai.FunctionParameters if gollmDef.Parameters == nil { return params, nil } // Convert the schema for OpenAI compatibility klog.V(2).Infof("Original schema for function %s: %+v", gollmDef.Name, gollmDef.Parameters) validatedSchema, err := convertSchemaForOpenAI(gollmDef.Parameters) if err != nil { return params, fmt.Errorf("schema conversion failed: %w", err) } klog.V(2).Infof("Converted schema for function %s: %+v", gollmDef.Name, validatedSchema) // Convert to raw schema bytes schemaBytes, err := cs.convertSchemaToBytes(validatedSchema, gollmDef.Name) if err != nil { return params, err } // Unmarshal into OpenAI parameters format if err := json.Unmarshal(schemaBytes, ¶ms); err != nil { return params, fmt.Errorf("failed to unmarshal schema: %w", err) } return params, nil } // openAISchema wraps a gollm Schema with OpenAI-specific marshaling behavior type openAISchema struct { *Schema } // MarshalJSON provides OpenAI-specific JSON marshaling that ensures object schemas have properties func (s openAISchema) MarshalJSON() ([]byte, error) { // Create a map to build the JSON representation result := make(map[string]interface{}) if s.Type != "" { result["type"] = s.Type } if s.Description != "" { result["description"] = s.Description } if len(s.Required) > 0 { result["required"] = s.Required } // For object types, always include properties (even if empty) to satisfy OpenAI if s.Type == TypeObject { if s.Properties != nil { result["properties"] = s.Properties } else { result["properties"] = make(map[string]*Schema) } } else if s.Properties != nil && len(s.Properties) > 0 { // For non-object types, only include properties if they exist and are non-empty result["properties"] = s.Properties } if s.Items != nil { result["items"] = s.Items } return json.Marshal(result) } // convertSchemaToBytes converts a validated schema to JSON bytes using OpenAI-specific marshaling func (cs *openAIChatSession) convertSchemaToBytes(schema *Schema, functionName string) ([]byte, error) { // Wrap the schema with OpenAI-specific marshaling behavior openAIWrapper := openAISchema{Schema: schema} bytes, err := json.Marshal(openAIWrapper) if err != nil { return nil, fmt.Errorf("failed to convert schema: %w", err) } klog.Infof("OpenAI schema for function %s: %s", functionName, string(bytes)) return bytes, nil } // newOpenAIClientFactory is the factory function for creating OpenAI clients. func newOpenAIClientFactory(ctx context.Context, opts ClientOptions) (Client, error) { return NewOpenAIClient(ctx, opts) } // addContentsToHistory processes and appends user messages to chat history func (cs *openAIChatSession) addContentsToHistory(contents []any) error { for _, content := range contents { switch c := content.(type) { case string: klog.V(2).Infof("Adding user message to history: %s", c) cs.history = append(cs.history, openai.UserMessage(c)) case FunctionCallResult: klog.V(2).Infof("Adding tool call result to history: Name=%s, ID=%s", c.Name, c.ID) // Marshal the result map into a JSON string for the message content resultJSON, err := json.Marshal(c.Result) if err != nil { klog.Errorf("Failed to marshal function call result: %v", err) return fmt.Errorf("failed to marshal function call result %q: %w", c.Name, err) } cs.history = append(cs.history, openai.ToolMessage(string(resultJSON), c.ID)) default: klog.Warningf("Unhandled content type: %T", content) return fmt.Errorf("unhandled content type: %T", content) } } return nil } // convertToolCallsToFunctionCalls converts OpenAI tool calls to gollm function calls func convertToolCallsToFunctionCalls(toolCalls []openai.ChatCompletionMessageToolCall) ([]FunctionCall, bool) { if len(toolCalls) == 0 { return nil, false } calls := make([]FunctionCall, 0, len(toolCalls)) for _, tc := range toolCalls { // Skip non-function tool calls if tc.Function.Name == "" { klog.V(2).Infof("Skipping non-function tool call ID: %s", tc.ID) continue } // Parse function arguments with error handling var args map[string]any if tc.Function.Arguments != "" { if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil { klog.V(2).Infof("Error unmarshalling function arguments for %s: %v", tc.Function.Name, err) args = make(map[string]any) } } else { args = make(map[string]any) } calls = append(calls, FunctionCall{ ID: tc.ID, Name: tc.Function.Name, Arguments: args, }) } return calls, len(calls) > 0 } // getOpenAIModel returns the appropriate model based on configuration and explicitly provided model name func getOpenAIModel(model string) string { // If explicit model is provided, use it if model != "" { klog.V(2).Infof("Using explicitly provided model: %s", model) return model } // Check configuration configModel := openAIModel if configModel != "" { klog.V(1).Infof("Using model from config: %s", configModel) return configModel } // Default model as fallback klog.V(2).Info("No model specified, defaulting to gpt-4.1") return "gpt-4.1" } ================================================ FILE: gollm/openai_response.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "context" "encoding/json" "errors" "fmt" "log" openai "github.com/openai/openai-go" "github.com/openai/openai-go/responses" "k8s.io/klog/v2" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) // Chat Session Implementation type openAIResponseChatSession struct { client openai.Client history responses.ResponseInputParam model string functionDefinitions []*FunctionDefinition // Stored in gollm format tools []responses.ToolUnionParam // Stored in OpenAI format // params to be intialized at the beginning of the session params responses.ResponseNewParams } // Ensure openAIChatSession implements the Chat interface. var _ Chat = (*openAIChatSession)(nil) // SetFunctionDefinitions stores the function definitions and converts them to OpenAI format. func (cs *openAIResponseChatSession) SetFunctionDefinitions(defs []*FunctionDefinition) error { cs.functionDefinitions = defs cs.tools = nil // Clear previous tools if len(defs) > 0 { cs.tools = make([]responses.ToolUnionParam, len(defs)) for i, gollmDef := range defs { klog.Infof("Processing function definition: %s", gollmDef.Name) // Process function parameters params, err := cs.convertFunctionParameters(gollmDef) if err != nil { return fmt.Errorf("failed to process parameters for function %s: %w", gollmDef.Name, err) } cs.tools[i] = responses.ToolUnionParam{ OfFunction: &responses.FunctionToolParam{ Name: gollmDef.Name, Description: openai.String(gollmDef.Description), Parameters: params, }, } } } klog.V(1).Infof("Set %d function definitions for OpenAI chat session", len(cs.functionDefinitions)) return nil } // Send sends the user message(s), appends to history, and gets the LLM response. func (cs *openAIResponseChatSession) Send(ctx context.Context, contents ...any) (ChatResponse, error) { klog.V(1).InfoS("openAIChatSession.Send called", "model", cs.model, "history_len", len(cs.history)) // TODO(droot): kubectl-ai agent uses SendStreaming instead of Send so deferred the implementation for now. return &openAIChatResponse{}, errors.ErrUnsupported } // SendStreaming sends the user message(s) and returns an iterator for the LLM response stream. func (cs *openAIResponseChatSession) SendStreaming(ctx context.Context, contents ...any) (ChatResponseIterator, error) { klog.V(1).InfoS("Starting OpenAI streaming request", "model", cs.model) // Process and append messages to history if err := cs.addContentsToHistory(contents); err != nil { return nil, err } // Prepare and send API request cs.params.Input = responses.ResponseNewParamsInputUnion{ OfInputItemList: cs.history, } cs.params.Tools = cs.tools klog.V(1).InfoS("Sending streaming request to OpenAI API", "model", cs.model, "messageCount", len(cs.params.Input.OfInputItemList), "toolCount", len(cs.params.Tools)) resp, err := cs.client.Responses.New(ctx, cs.params) if err == nil { for _, output := range resp.Output { switch output.AsAny().(type) { case responses.ResponseFunctionToolCall: fc := output.AsFunctionCall() log.Printf("Inspected function call item: %+v", fc) fpP := fc.ToParam() cs.history = append(cs.history, responses.ResponseInputItemUnionParam{ OfFunctionCall: &fpP, }) case responses.ResponseReasoningItem: reason := output.AsReasoning() log.Printf("Inspected Reasoning item: %+v", reason) reasonParam := reason.ToParam() cs.history = append(cs.history, responses.ResponseInputItemUnionParam{ OfReasoning: &reasonParam, }) case responses.ResponseOutputMessage: msg := output.AsMessage() log.Printf("Inspected Output Message: %+v", msg.Content[0].Text) cs.history = append(cs.history, responses.ResponseInputItemParamOfOutputMessage([]responses.ResponseOutputMessageContentUnionParam{ { OfOutputText: &responses.ResponseOutputTextParam{ Annotations: []responses.ResponseOutputTextAnnotationUnionParam{}, Text: msg.Content[0].Text, Type: "output_text", }, }, }, msg.ID, msg.Status), ) default: log.Println("no variant present", output) } } } return singletonChatResponseIterator(&openAIResponseChatResponse{resp: resp}), err } // IsRetryableError determines if an error from the OpenAI API should be retried. func (cs *openAIResponseChatSession) IsRetryableError(err error) bool { if err == nil { return false } return DefaultIsRetryableError(err) } func (cs *openAIResponseChatSession) Initialize(messages []*api.Message) error { klog.Warning("chat history persistence is not supported for provider 'openai', using in-memory chat history") return nil } // Helper structs for ChatResponse interface type openAIResponseChatResponse struct { resp *responses.Response } var _ ChatResponse = (*openAIResponseChatResponse)(nil) func (r *openAIResponseChatResponse) UsageMetadata() any { return nil } func (r *openAIResponseChatResponse) Candidates() []Candidate { if r.resp == nil { return nil } var candidates []Candidate for _, output := range r.resp.Output { switch output.AsAny().(type) { case responses.ResponseFunctionToolCall, responses.ResponseOutputMessage: candidates = append(candidates, &openAIResponseCandidate{ candidate: &output, }) default: // skip reasoning messages because agentic loop doesn't know // how to handle them yet. } } return candidates } type openAIResponseCandidate struct { candidate *responses.ResponseOutputItemUnion } var _ Candidate = (*openAIResponseCandidate)(nil) func (c *openAIResponseCandidate) Parts() []Part { if c.candidate == nil { return nil } // OpenAI message can have Content AND ToolCalls var parts []Part output := c.candidate switch output.AsAny().(type) { case responses.ResponseFunctionToolCall: fc := output.AsFunctionCall() toolCall, err := convertResponseToolCallToFunctionCall(fc) if err != nil { // } parts = append(parts, &openAIResponsePart{ toolCall: toolCall, }) case responses.ResponseReasoningItem: reason := output.AsReasoning() log.Printf("Inspected Reasoning item: %+v", reason) case responses.ResponseOutputMessage: msg := output.AsMessage() parts = append(parts, &openAIResponsePart{ content: msg.Content[0].AsOutputText().Text, }) default: log.Println("no variant present", output) } return parts } // String provides a simple string representation for logging/debugging. func (c *openAIResponseCandidate) String() string { return fmt.Sprintf("%+v", c.candidate) } type openAIResponsePart struct { content string toolCall FunctionCall } var _ Part = (*openAIResponsePart)(nil) func (p *openAIResponsePart) AsText() (string, bool) { return p.content, p.content != "" } func (p *openAIResponsePart) AsFunctionCalls() ([]FunctionCall, bool) { return []FunctionCall{p.toolCall}, p.content == "" } // convertFunctionParameters handles the conversion of gollm parameters to OpenAI format func (cs *openAIResponseChatSession) convertFunctionParameters(gollmDef *FunctionDefinition) (openai.FunctionParameters, error) { var params openai.FunctionParameters if gollmDef.Parameters == nil { return params, nil } // Convert the schema for OpenAI compatibility klog.V(2).Infof("Original schema for function %s: %+v", gollmDef.Name, gollmDef.Parameters) validatedSchema, err := convertSchemaForOpenAI(gollmDef.Parameters) if err != nil { return params, fmt.Errorf("schema conversion failed: %w", err) } klog.V(2).Infof("Converted schema for function %s: %+v", gollmDef.Name, validatedSchema) // Convert to raw schema bytes schemaBytes, err := cs.convertSchemaToBytes(validatedSchema, gollmDef.Name) if err != nil { return params, err } // Unmarshal into OpenAI parameters format if err := json.Unmarshal(schemaBytes, ¶ms); err != nil { return params, fmt.Errorf("failed to unmarshal schema: %w", err) } return params, nil } // convertSchemaToBytes converts a validated schema to JSON bytes using OpenAI-specific marshaling func (cs *openAIResponseChatSession) convertSchemaToBytes(schema *Schema, functionName string) ([]byte, error) { // Wrap the schema with OpenAI-specific marshaling behavior openAIWrapper := openAISchema{Schema: schema} bytes, err := json.Marshal(openAIWrapper) if err != nil { return nil, fmt.Errorf("failed to convert schema: %w", err) } klog.Infof("OpenAI schema for function %s: %s", functionName, string(bytes)) return bytes, nil } // addContentsToHistory processes and appends user messages to chat history func (cs *openAIResponseChatSession) addContentsToHistory(contents []any) error { for _, content := range contents { switch c := content.(type) { case string: klog.V(2).Infof("Adding user message to history: %s", c) cs.history = append(cs.history, responses.ResponseInputItemUnionParam{ OfMessage: &responses.EasyInputMessageParam{ Content: responses.EasyInputMessageContentUnionParam{ OfString: openai.String(c), }, Role: responses.EasyInputMessageRoleUser, }, }) case FunctionCallResult: klog.V(2).Infof("Adding tool call result to history: Name=%s, ID=%s", c.Name, c.ID) // Marshal the result map into a JSON string for the message content resultJSON, err := json.Marshal(c.Result) if err != nil { klog.Errorf("Failed to marshal function call result: %v", err) return fmt.Errorf("failed to marshal function call result %q: %w", c.Name, err) } // cs.history = append(cs.history, openai.ToolMessage(string(resultJSON), c.ID)) cs.history = append(cs.history, responses.ResponseInputItemParamOfFunctionCallOutput(c.ID, string(resultJSON))) default: klog.Warningf("Unhandled content type: %T", content) return fmt.Errorf("unhandled content type: %T", content) } } return nil } // convertToolCallsToFunctionCalls converts OpenAI tool calls to gollm function calls func convertResponseToolCallToFunctionCall(responseToolCall responses.ResponseFunctionToolCall) (FunctionCall, error) { fc := FunctionCall{} // Skip non-function tool calls if responseToolCall.Name == "" { klog.V(2).Infof("Skipping non-function tool call ID: %s", responseToolCall.ID) return fc, fmt.Errorf("missing name %v", responseToolCall) } fc.Name = responseToolCall.Name // Parse function arguments with error handling var args map[string]any if responseToolCall.Arguments != "" { if err := json.Unmarshal([]byte(responseToolCall.Arguments), &args); err != nil { klog.V(2).Infof("Error unmarshalling function arguments for %s: %v", fc.Name, err) args = make(map[string]any) } } else { args = make(map[string]any) } return FunctionCall{ ID: responseToolCall.CallID, Name: responseToolCall.Name, Arguments: args, }, nil } ================================================ FILE: gollm/openai_test.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "encoding/json" "testing" "github.com/openai/openai-go" ) func TestConvertSchemaForOpenAI(t *testing.T) { tests := []struct { name string inputSchema *Schema expectedType SchemaType expectedError bool validateResult func(t *testing.T, result *Schema) }{ // Core logic tests { name: "nil schema", inputSchema: nil, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Properties == nil { t.Error("expected properties map to be initialized") } if len(result.Properties) != 0 { t.Error("expected empty properties map") } }, }, { name: "simple string schema", inputSchema: &Schema{ Type: TypeString, Description: "A simple string", }, expectedType: TypeString, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Description != "A simple string" { t.Errorf("expected description 'A simple string', got %q", result.Description) } }, }, { name: "simple number schema", inputSchema: &Schema{ Type: TypeNumber, }, expectedType: TypeNumber, expectedError: false, }, { name: "integer schema converted to number", inputSchema: &Schema{ Type: TypeInteger, Description: "An integer value", }, expectedType: TypeNumber, // OpenAI prefers number for integers expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Description != "An integer value" { t.Errorf("expected description preserved") } }, }, { name: "boolean schema", inputSchema: &Schema{ Type: TypeBoolean, }, expectedType: TypeBoolean, expectedError: false, }, { name: "empty type defaults to object", inputSchema: &Schema{ Description: "No type specified", }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Properties == nil { t.Error("expected properties map to be initialized") } }, }, { name: "unknown type defaults to object", inputSchema: &Schema{ Type: "unknown", }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Properties == nil { t.Error("expected properties map to be initialized") } }, }, { name: "object schema with properties", inputSchema: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "name": {Type: TypeString, Description: "User name"}, "age": {Type: TypeInteger, Description: "User age"}, }, Required: []string{"name"}, }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if len(result.Properties) != 2 { t.Errorf("expected 2 properties, got %d", len(result.Properties)) } if result.Properties["name"].Type != TypeString { t.Error("expected name property to be string") } // Age should be converted from integer to number if result.Properties["age"].Type != TypeNumber { t.Error("expected age property to be converted to number") } if len(result.Required) != 1 || result.Required[0] != "name" { t.Error("expected required fields to be preserved") } }, }, { name: "object schema without properties", inputSchema: &Schema{ Type: TypeObject, }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Properties == nil { t.Error("expected properties map to be initialized") } if len(result.Properties) != 0 { t.Error("expected empty properties map") } }, }, { name: "array schema with string items", inputSchema: &Schema{ Type: TypeArray, Items: &Schema{Type: TypeString}, }, expectedType: TypeArray, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Items == nil { t.Error("expected items schema to be present") } if result.Items.Type != TypeString { t.Error("expected items to be string type") } }, }, { name: "array schema with integer items (converted to number)", inputSchema: &Schema{ Type: TypeArray, Items: &Schema{Type: TypeInteger}, }, expectedType: TypeArray, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Items == nil { t.Error("expected items schema to be present") } if result.Items.Type != TypeNumber { t.Error("expected items to be converted to number type") } }, }, { name: "array schema without items (defaults to string)", inputSchema: &Schema{ Type: TypeArray, }, expectedType: TypeArray, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Items == nil { t.Error("expected items schema to be defaulted") } if result.Items.Type != TypeString { t.Error("expected default items to be string type") } }, }, { name: "nested object in array", inputSchema: &Schema{ Type: TypeArray, Items: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "id": {Type: TypeInteger}, "name": {Type: TypeString}, }, }, }, expectedType: TypeArray, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if result.Items == nil { t.Error("expected items schema to be present") } if result.Items.Type != TypeObject { t.Error("expected items to be object type") } if result.Items.Properties["id"].Type != TypeNumber { t.Error("expected nested integer to be converted to number") } if result.Items.Properties["name"].Type != TypeString { t.Error("expected nested string to remain string") } }, }, // Built-in tool schema tests { name: "kubectl tool schema", inputSchema: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "command": { Type: TypeString, Description: "The complete kubectl command to execute", }, "modifies_resource": { Type: TypeString, Description: "Whether the command modifies a kubernetes resource", }, }, }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if len(result.Properties) != 2 { t.Errorf("expected 2 properties, got %d", len(result.Properties)) } if result.Properties["command"].Type != TypeString { t.Error("expected command property to be string") } if result.Properties["modifies_resource"].Type != TypeString { t.Error("expected modifies_resource property to be string") } // Properties should be initialized if result.Properties == nil { t.Error("expected properties to be initialized") } }, }, { name: "bash tool schema", inputSchema: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "command": { Type: TypeString, Description: "The bash command to execute", }, "modifies_resource": { Type: TypeString, Description: "Whether the command modifies a kubernetes resource", }, }, }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { if len(result.Properties) != 2 { t.Errorf("expected 2 properties, got %d", len(result.Properties)) } // All string properties should remain strings if result.Properties["command"].Type != TypeString { t.Error("expected command property to remain string") } if result.Properties["modifies_resource"].Type != TypeString { t.Error("expected modifies_resource property to remain string") } }, }, { name: "mcp tool schema with complex nested structure", inputSchema: &Schema{ Type: TypeObject, Properties: map[string]*Schema{ "server_name": { Type: TypeString, Description: "Name of the MCP server", }, "method": { Type: TypeString, Description: "MCP method name", }, "params": { Type: TypeObject, Properties: map[string]*Schema{ "query": {Type: TypeString}, "limit": {Type: TypeInteger}, // Should convert to number }, }, }, Required: []string{"server_name", "method"}, }, expectedType: TypeObject, expectedError: false, validateResult: func(t *testing.T, result *Schema) { // Check top-level properties if len(result.Properties) != 3 { t.Errorf("expected 3 properties, got %d", len(result.Properties)) } // Check nested object conversion params := result.Properties["params"] if params.Type != TypeObject { t.Error("expected params to be object type") } if params.Properties == nil { t.Error("expected params properties to be initialized") } // Check nested integer conversion if params.Properties["limit"].Type != TypeNumber { t.Error("expected nested limit property to be converted to number") } // Check required fields preservation if len(result.Required) != 2 { t.Error("expected required fields to be preserved") } }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result, err := convertSchemaForOpenAI(tt.inputSchema) if tt.expectedError { if err == nil { t.Error("expected error but got none") } return } if err != nil { t.Errorf("unexpected error: %v", err) return } if result == nil { t.Error("expected non-nil result") return } if result.Type != tt.expectedType { t.Errorf("expected type %q, got %q", tt.expectedType, result.Type) } // Run custom validation if provided if tt.validateResult != nil { tt.validateResult(t, result) } }) } } // TestConvertSchemaToBytes tests the JSON-level fix for the omitempty issue func TestConvertSchemaToBytes(t *testing.T) { session := &openAIChatSession{} // Test case: Object schema with empty properties map (which gets omitted by omitempty) schema := &Schema{ Type: TypeObject, Properties: make(map[string]*Schema), // Empty map gets omitted by omitempty } bytes, err := session.convertSchemaToBytes(schema, "test_function") if err != nil { t.Errorf("unexpected error: %v", err) return } // Parse the JSON to verify it has properties field var schemaMap map[string]any if err := json.Unmarshal(bytes, &schemaMap); err != nil { t.Errorf("failed to unmarshal schema: %v", err) return } // Verify the schema has type: object if schemaType, ok := schemaMap["type"].(string); !ok || schemaType != "object" { t.Errorf("expected type 'object', got %v", schemaMap["type"]) } // Verify the schema has properties field (even if empty) if _, hasProperties := schemaMap["properties"]; !hasProperties { t.Error("expected properties field to be present in JSON, but it was missing") } // Verify properties is an empty object if props, ok := schemaMap["properties"].(map[string]any); !ok { t.Error("expected properties to be an object") } else if len(props) != 0 { t.Errorf("expected empty properties object, got %v", props) } } // TestConvertToolCallsToFunctionCalls tests the tool call conversion logic func TestConvertToolCallsToFunctionCalls(t *testing.T) { tests := []struct { name string toolCalls []openai.ChatCompletionMessageToolCall expectedCount int expectedResult bool validateCalls func(t *testing.T, calls []FunctionCall) }{ { name: "empty tool calls", toolCalls: []openai.ChatCompletionMessageToolCall{}, expectedCount: 0, expectedResult: false, }, { name: "nil tool calls", toolCalls: nil, expectedCount: 0, expectedResult: false, }, { name: "single valid tool call", toolCalls: []openai.ChatCompletionMessageToolCall{ { ID: "call_123", Function: openai.ChatCompletionMessageToolCallFunction{ Name: "kubectl", Arguments: `{"command":"kubectl get pods --namespace=app-dev01","modifies_resource":"no"}`, }, }, }, expectedCount: 1, expectedResult: true, validateCalls: func(t *testing.T, calls []FunctionCall) { if calls[0].ID != "call_123" { t.Errorf("expected ID 'call_123', got %s", calls[0].ID) } if calls[0].Name != "kubectl" { t.Errorf("expected Name 'kubectl', got %s", calls[0].Name) } if calls[0].Arguments["command"] != "kubectl get pods --namespace=app-dev01" { t.Errorf("expected command argument, got %v", calls[0].Arguments["command"]) } if calls[0].Arguments["modifies_resource"] != "no" { t.Errorf("expected modifies_resource argument, got %v", calls[0].Arguments["modifies_resource"]) } }, }, { name: "tool call with empty function name", toolCalls: []openai.ChatCompletionMessageToolCall{ { ID: "call_456", Function: openai.ChatCompletionMessageToolCallFunction{ Name: "", Arguments: `{"command":"kubectl get pods"}`, }, }, }, expectedCount: 0, expectedResult: false, }, { name: "tool call with invalid JSON arguments", toolCalls: []openai.ChatCompletionMessageToolCall{ { ID: "call_789", Function: openai.ChatCompletionMessageToolCallFunction{ Name: "kubectl", Arguments: `{"command":"kubectl get pods", invalid json}`, }, }, }, expectedCount: 1, expectedResult: true, validateCalls: func(t *testing.T, calls []FunctionCall) { if calls[0].ID != "call_789" { t.Errorf("expected ID 'call_789', got %s", calls[0].ID) } if calls[0].Name != "kubectl" { t.Errorf("expected Name 'kubectl', got %s", calls[0].Name) } // Arguments should be empty due to parsing error if len(calls[0].Arguments) != 0 { t.Errorf("expected empty arguments due to parse error, got %v", calls[0].Arguments) } }, }, { name: "tool call with empty arguments", toolCalls: []openai.ChatCompletionMessageToolCall{ { ID: "call_empty", Function: openai.ChatCompletionMessageToolCallFunction{ Name: "kubectl", Arguments: "", }, }, }, expectedCount: 1, expectedResult: true, validateCalls: func(t *testing.T, calls []FunctionCall) { if calls[0].ID != "call_empty" { t.Errorf("expected ID 'call_empty', got %s", calls[0].ID) } if calls[0].Name != "kubectl" { t.Errorf("expected Name 'kubectl', got %s", calls[0].Name) } // Arguments should be empty but not nil if calls[0].Arguments == nil { t.Error("expected non-nil arguments map") } if len(calls[0].Arguments) != 0 { t.Errorf("expected empty arguments, got %v", calls[0].Arguments) } }, }, { name: "multiple tool calls with reasoning model pattern", toolCalls: []openai.ChatCompletionMessageToolCall{ { ID: "call_1", Function: openai.ChatCompletionMessageToolCallFunction{ Name: "kubectl", Arguments: `{"command":"kubectl get pods --namespace=app-dev01\nkubectl get pods --namespace=app-dev02","modifies_resource":"no"}`, }, }, { ID: "call_2", Function: openai.ChatCompletionMessageToolCallFunction{ Name: "bash", Arguments: `{"command":"echo 'test'","modifies_resource":"no"}`, }, }, }, expectedCount: 2, expectedResult: true, validateCalls: func(t *testing.T, calls []FunctionCall) { if len(calls) != 2 { t.Errorf("expected 2 calls, got %d", len(calls)) } // Check first call if calls[0].Name != "kubectl" { t.Errorf("expected first call to be 'kubectl', got %s", calls[0].Name) } if calls[0].Arguments["command"] != "kubectl get pods --namespace=app-dev01\nkubectl get pods --namespace=app-dev02" { t.Errorf("expected multi-line command, got %v", calls[0].Arguments["command"]) } // Check second call if calls[1].Name != "bash" { t.Errorf("expected second call to be 'bash', got %s", calls[1].Name) } if calls[1].Arguments["command"] != "echo 'test'" { t.Errorf("expected echo command, got %v", calls[1].Arguments["command"]) } }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { calls, ok := convertToolCallsToFunctionCalls(tt.toolCalls) if ok != tt.expectedResult { t.Errorf("expected result %v, got %v", tt.expectedResult, ok) } if len(calls) != tt.expectedCount { t.Errorf("expected %d calls, got %d", tt.expectedCount, len(calls)) } if tt.validateCalls != nil && len(calls) > 0 { tt.validateCalls(t, calls) } }) } } ================================================ FILE: gollm/persist.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm // We define some standard structs to allow for persistence of the LLM requests and responses. // This lets us store the history of the conversation for later analysis. type RecordCompletionResponse struct { Text string `json:"text"` Raw any `json:"raw"` } type RecordChatResponse struct { // TODO: Structured data? Raw any `json:"raw"` } ================================================ FILE: gollm/schema.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm import ( "reflect" "strings" "k8s.io/klog/v2" ) // BuildSchemaFor will build a schema for the given golang type. // Because this does not have description populated, it is more useful for the response schema than tools/functions. func BuildSchemaFor(t reflect.Type) *Schema { out := &Schema{} switch t.Kind() { case reflect.String: out.Type = TypeString case reflect.Bool: out.Type = TypeBoolean case reflect.Int: out.Type = TypeInteger case reflect.Struct: out.Type = TypeObject out.Properties = make(map[string]*Schema) numFields := t.NumField() required := []string{} for i := 0; i < numFields; i++ { field := t.Field(i) jsonTag := field.Tag.Get("json") if jsonTag == "" { continue } if strings.HasSuffix(jsonTag, ",omitempty") { jsonTag = strings.TrimSuffix(jsonTag, ",omitempty") } else { required = append(required, jsonTag) } fieldType := field.Type fieldSchema := BuildSchemaFor(fieldType) out.Properties[jsonTag] = fieldSchema } if len(required) != 0 { out.Required = required } case reflect.Slice: out.Type = TypeArray out.Items = BuildSchemaFor(t.Elem()) default: klog.Fatalf("unhandled kind %v", t.Kind()) } return out } ================================================ FILE: gollm/shims.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package gollm func singletonChatResponseIterator(response ChatResponse) ChatResponseIterator { return func(yield func(ChatResponse, error) bool) { if !yield(response, nil) { return } } } ================================================ FILE: images/kubectl-ai/Dockerfile ================================================ ARG GO_VERSION="1.24.3" ARG GCLOUD_CLI_VERSION="530.0.0-stable" FROM golang:${GO_VERSION}-bookworm AS builder WORKDIR /src COPY go.mod go.sum ./ COPY gollm/ ./gollm/ RUN go mod download COPY cmd/ ./cmd/ COPY pkg/ ./pkg/ RUN CGO_ENABLED=0 go build -o kubectl-ai ./cmd/ FROM debian:bookworm-slim AS kubectl-tool ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update && \ apt-get install -y --no-install-recommends curl ca-certificates && \ mkdir -p /opt/tools/kubectl/bin/ && \ curl -v -L "https://dl.k8s.io/release/v1.33.0/bin/linux/amd64/kubectl" -o /opt/tools/kubectl/bin/kubectl && \ chmod +x /opt/tools/kubectl/bin/kubectl && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* FROM gcr.io/google.com/cloudsdktool/google-cloud-cli:${GCLOUD_CLI_VERSION} AS runtime RUN apt-get update -y && \ apt-get install -y apt-transport-https ca-certificates gnupg curl ca-certificates && \ curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \ echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee /etc/apt/sources.list.d/google-cloud-sdk.list && \ apt-get update -y && \ apt-get install -y google-cloud-cli-gke-gcloud-auth-plugin && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* COPY --from=builder /src/kubectl-ai /bin/kubectl-ai COPY --from=kubectl-tool /opt/tools/kubectl/ /opt/tools/kubectl/ RUN ln -sf /opt/tools/kubectl/bin/kubectl /bin/kubectl # Copy the custom tool configurations into the runtime image. COPY docs/tool-samples /etc/kubectl-ai/tools/ ENTRYPOINT [ "/bin/kubectl-ai" ] ================================================ FILE: install.sh ================================================ #!/usr/bin/env bash set -euo pipefail # Check for required commands for cmd in curl tar; do if ! command -v "$cmd" >/dev/null 2>&1; then echo "Error: $cmd is not installed. Please install $cmd to proceed." exit 1 fi done # Set the insecure SSL argument INSECURE_ARG="" if [ -n "${INSECURE:-}" ]; then INSECURE_ARG="--insecure" fi REPO="GoogleCloudPlatform/kubectl-ai" BINARY="kubectl-ai" # Detect OS sysOS="$(uname | tr '[:upper:]' '[:lower:]')" case "$sysOS" in linux) OS="Linux" ;; darwin) OS="Darwin" ;; *) echo "If you are on Windows or another unsupported OS, please follow the manual installation instructions at:" echo "https://github.com/GoogleCloudPlatform/kubectl-ai#manual-installation-linux-macos-and-windows" exit 1 ;; esac # Detect NixOS nixos_check="$(grep "ID=nixos" /etc/os-release 2>/dev/null || echo "no-match")" case "$nixos_check" in *nixos*) echo "NixOS detected, please follow the manual installation instructions at:" echo "https://github.com/GoogleCloudPlatform/kubectl-ai#install-on-nixos" exit 1 ;; esac # Detect ARCH ARCH="$(uname -m)" case "$ARCH" in x86_64|amd64) ARCH="x86_64" ;; arm64|aarch64) ARCH="arm64" ;; *) echo "If you are on an unsupported architecture, please follow the manual installation instructions at:" echo "https://github.com/GoogleCloudPlatform/kubectl-ai#manual-installation-linux-macos-and-windows" exit 1 ;; esac # Get latest version tag from GitHub API, Use GITHUB_TOKEN if available to avoid potential rate limit if [ -n "${GITHUB_TOKEN:-}" ]; then auth_hdr="Authorization: token $GITHUB_TOKEN" else auth_hdr="" fi if [ -n "${INSECURE:-}" ]; then echo "⚠️ SECURITY WARNING: INSECURE is set, SSL certificate validation will be skipped!" echo " This makes you vulnerable to man-in-the-middle attacks and other security risks." echo " Only proceed if you understand the security implications and trust your network." echo "" echo " Continue with unsafe download? (yes/no)" read -r response case "$response" in [yY][eE][sS]|[yY]) echo "Proceeding with insecure connection..." ;; *) echo "Installation aborted for security reasons." exit 1 ;; esac fi LATEST_TAG=$(curl $INSECURE_ARG -s -H "$auth_hdr" \ "https://api.github.com/repos/$REPO/releases/latest" \ | sed -n 's/.*"tag_name": *"\([^"]*\)".*/\1/p') if [ -z "$LATEST_TAG" ]; then echo "Failed to fetch latest release tag." exit 1 fi # Compose download URL TARBALL="kubectl-ai_${OS}_${ARCH}.tar.gz" URL="https://github.com/$REPO/releases/download/$LATEST_TAG/$TARBALL" # Create temp dir and set cleanup trap TEMP_DIR="$(mktemp -d)" cleanup() { if [ -n "${TEMP_DIR:-}" ] && [ -d "$TEMP_DIR" ]; then rm -rf "$TEMP_DIR" fi } trap cleanup EXIT INT TERM # Download and extract in temp dir; install from there ( cd "$TEMP_DIR" echo "Downloading $URL ..." CURL_FLAGS="-fSL --retry 3" if [ -n "${INSECURE:-}" ]; then echo "⚠️ SSL certificate validation will be skipped for this download." fi curl $INSECURE_ARG $CURL_FLAGS "$URL" -o "$TARBALL" tar --no-same-owner -xzf "$TARBALL" if [ ! -f "$BINARY" ]; then echo "Error: expected binary '$BINARY' not found after extraction." exit 1 fi INSTALL_DIR="${INSTALL_DIR:-/usr/local/bin}" echo "Installing $BINARY to $INSTALL_DIR (may require sudo)..." sudo install -m 0755 "$BINARY" "$INSTALL_DIR/" ) echo "✅ $BINARY installed successfully! Run '$BINARY --help' to get started." ================================================ FILE: internal/mocks/agent_mock.go ================================================ // Code generated by MockGen. DO NOT EDIT. // Source: github.com/GoogleCloudPlatform/kubectl-ai/pkg/api (interfaces: ChatMessageStore) // // Generated by this command: // // mockgen -destination=agent_mock.go -package=mocks github.com/GoogleCloudPlatform/kubectl-ai/pkg/api ChatMessageStore // // Package mocks is a generated GoMock package. package mocks import ( reflect "reflect" api "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" gomock "go.uber.org/mock/gomock" ) // MockChatMessageStore is a mock of ChatMessageStore interface. type MockChatMessageStore struct { ctrl *gomock.Controller recorder *MockChatMessageStoreMockRecorder isgomock struct{} } // MockChatMessageStoreMockRecorder is the mock recorder for MockChatMessageStore. type MockChatMessageStoreMockRecorder struct { mock *MockChatMessageStore } // NewMockChatMessageStore creates a new mock instance. func NewMockChatMessageStore(ctrl *gomock.Controller) *MockChatMessageStore { mock := &MockChatMessageStore{ctrl: ctrl} mock.recorder = &MockChatMessageStoreMockRecorder{mock} return mock } // EXPECT returns an object that allows the caller to indicate expected use. func (m *MockChatMessageStore) EXPECT() *MockChatMessageStoreMockRecorder { return m.recorder } // AddChatMessage mocks base method. func (m *MockChatMessageStore) AddChatMessage(record *api.Message) error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "AddChatMessage", record) ret0, _ := ret[0].(error) return ret0 } // AddChatMessage indicates an expected call of AddChatMessage. func (mr *MockChatMessageStoreMockRecorder) AddChatMessage(record any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddChatMessage", reflect.TypeOf((*MockChatMessageStore)(nil).AddChatMessage), record) } // ChatMessages mocks base method. func (m *MockChatMessageStore) ChatMessages() []*api.Message { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "ChatMessages") ret0, _ := ret[0].([]*api.Message) return ret0 } // ChatMessages indicates an expected call of ChatMessages. func (mr *MockChatMessageStoreMockRecorder) ChatMessages() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ChatMessages", reflect.TypeOf((*MockChatMessageStore)(nil).ChatMessages)) } // ClearChatMessages mocks base method. func (m *MockChatMessageStore) ClearChatMessages() error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "ClearChatMessages") ret0, _ := ret[0].(error) return ret0 } // ClearChatMessages indicates an expected call of ClearChatMessages. func (mr *MockChatMessageStoreMockRecorder) ClearChatMessages() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ClearChatMessages", reflect.TypeOf((*MockChatMessageStore)(nil).ClearChatMessages)) } // SetChatMessages mocks base method. func (m *MockChatMessageStore) SetChatMessages(newHistory []*api.Message) error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "SetChatMessages", newHistory) ret0, _ := ret[0].(error) return ret0 } // SetChatMessages indicates an expected call of SetChatMessages. func (mr *MockChatMessageStoreMockRecorder) SetChatMessages(newHistory any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetChatMessages", reflect.TypeOf((*MockChatMessageStore)(nil).SetChatMessages), newHistory) } ================================================ FILE: internal/mocks/generate.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Package mocks holds go:generate directives for gomock. package mocks // Generate gomock types for external interfaces we depend on. // NOTE: run `go generate ./...` from repo root to (re)create mocks. // Requires: go install go.uber.org/mock/mockgen@latest // gollm interfaces // - Client, Chat // tools interface // - Tool //go:generate mockgen -destination=gollm_mock.go -package=mocks github.com/GoogleCloudPlatform/kubectl-ai/gollm Client,Chat //go:generate mockgen -destination=tools_mock.go -package=mocks github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools Tool //go:generate mockgen -destination=agent_mock.go -package=mocks github.com/GoogleCloudPlatform/kubectl-ai/pkg/api ChatMessageStore ================================================ FILE: internal/mocks/gollm_mock.go ================================================ // Code generated by MockGen. DO NOT EDIT. // Source: github.com/GoogleCloudPlatform/kubectl-ai/gollm (interfaces: Client,Chat) // // Generated by this command: // // mockgen -destination=gollm_mock.go -package=mocks github.com/GoogleCloudPlatform/kubectl-ai/gollm Client,Chat // // Package mocks is a generated GoMock package. package mocks import ( context "context" reflect "reflect" gollm "github.com/GoogleCloudPlatform/kubectl-ai/gollm" api "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" gomock "go.uber.org/mock/gomock" ) // MockClient is a mock of Client interface. type MockClient struct { ctrl *gomock.Controller recorder *MockClientMockRecorder isgomock struct{} } // MockClientMockRecorder is the mock recorder for MockClient. type MockClientMockRecorder struct { mock *MockClient } // NewMockClient creates a new mock instance. func NewMockClient(ctrl *gomock.Controller) *MockClient { mock := &MockClient{ctrl: ctrl} mock.recorder = &MockClientMockRecorder{mock} return mock } // EXPECT returns an object that allows the caller to indicate expected use. func (m *MockClient) EXPECT() *MockClientMockRecorder { return m.recorder } // Close mocks base method. func (m *MockClient) Close() error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Close") ret0, _ := ret[0].(error) return ret0 } // Close indicates an expected call of Close. func (mr *MockClientMockRecorder) Close() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockClient)(nil).Close)) } // GenerateCompletion mocks base method. func (m *MockClient) GenerateCompletion(ctx context.Context, req *gollm.CompletionRequest) (gollm.CompletionResponse, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GenerateCompletion", ctx, req) ret0, _ := ret[0].(gollm.CompletionResponse) ret1, _ := ret[1].(error) return ret0, ret1 } // GenerateCompletion indicates an expected call of GenerateCompletion. func (mr *MockClientMockRecorder) GenerateCompletion(ctx, req any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateCompletion", reflect.TypeOf((*MockClient)(nil).GenerateCompletion), ctx, req) } // ListModels mocks base method. func (m *MockClient) ListModels(ctx context.Context) ([]string, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "ListModels", ctx) ret0, _ := ret[0].([]string) ret1, _ := ret[1].(error) return ret0, ret1 } // ListModels indicates an expected call of ListModels. func (mr *MockClientMockRecorder) ListModels(ctx any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListModels", reflect.TypeOf((*MockClient)(nil).ListModels), ctx) } // SetResponseSchema mocks base method. func (m *MockClient) SetResponseSchema(schema *gollm.Schema) error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "SetResponseSchema", schema) ret0, _ := ret[0].(error) return ret0 } // SetResponseSchema indicates an expected call of SetResponseSchema. func (mr *MockClientMockRecorder) SetResponseSchema(schema any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetResponseSchema", reflect.TypeOf((*MockClient)(nil).SetResponseSchema), schema) } // StartChat mocks base method. func (m *MockClient) StartChat(systemPrompt, model string) gollm.Chat { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "StartChat", systemPrompt, model) ret0, _ := ret[0].(gollm.Chat) return ret0 } // StartChat indicates an expected call of StartChat. func (mr *MockClientMockRecorder) StartChat(systemPrompt, model any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartChat", reflect.TypeOf((*MockClient)(nil).StartChat), systemPrompt, model) } // MockChat is a mock of Chat interface. type MockChat struct { ctrl *gomock.Controller recorder *MockChatMockRecorder isgomock struct{} } // MockChatMockRecorder is the mock recorder for MockChat. type MockChatMockRecorder struct { mock *MockChat } // NewMockChat creates a new mock instance. func NewMockChat(ctrl *gomock.Controller) *MockChat { mock := &MockChat{ctrl: ctrl} mock.recorder = &MockChatMockRecorder{mock} return mock } // EXPECT returns an object that allows the caller to indicate expected use. func (m *MockChat) EXPECT() *MockChatMockRecorder { return m.recorder } // Initialize mocks base method. func (m *MockChat) Initialize(messages []*api.Message) error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Initialize", messages) ret0, _ := ret[0].(error) return ret0 } // Initialize indicates an expected call of Initialize. func (mr *MockChatMockRecorder) Initialize(messages any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Initialize", reflect.TypeOf((*MockChat)(nil).Initialize), messages) } // IsRetryableError mocks base method. func (m *MockChat) IsRetryableError(arg0 error) bool { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "IsRetryableError", arg0) ret0, _ := ret[0].(bool) return ret0 } // IsRetryableError indicates an expected call of IsRetryableError. func (mr *MockChatMockRecorder) IsRetryableError(arg0 any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsRetryableError", reflect.TypeOf((*MockChat)(nil).IsRetryableError), arg0) } // Send mocks base method. func (m *MockChat) Send(ctx context.Context, contents ...any) (gollm.ChatResponse, error) { m.ctrl.T.Helper() varargs := []any{ctx} for _, a := range contents { varargs = append(varargs, a) } ret := m.ctrl.Call(m, "Send", varargs...) ret0, _ := ret[0].(gollm.ChatResponse) ret1, _ := ret[1].(error) return ret0, ret1 } // Send indicates an expected call of Send. func (mr *MockChatMockRecorder) Send(ctx any, contents ...any) *gomock.Call { mr.mock.ctrl.T.Helper() varargs := append([]any{ctx}, contents...) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MockChat)(nil).Send), varargs...) } // SendStreaming mocks base method. func (m *MockChat) SendStreaming(ctx context.Context, contents ...any) (gollm.ChatResponseIterator, error) { m.ctrl.T.Helper() varargs := []any{ctx} for _, a := range contents { varargs = append(varargs, a) } ret := m.ctrl.Call(m, "SendStreaming", varargs...) ret0, _ := ret[0].(gollm.ChatResponseIterator) ret1, _ := ret[1].(error) return ret0, ret1 } // SendStreaming indicates an expected call of SendStreaming. func (mr *MockChatMockRecorder) SendStreaming(ctx any, contents ...any) *gomock.Call { mr.mock.ctrl.T.Helper() varargs := append([]any{ctx}, contents...) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SendStreaming", reflect.TypeOf((*MockChat)(nil).SendStreaming), varargs...) } // SetFunctionDefinitions mocks base method. func (m *MockChat) SetFunctionDefinitions(functionDefinitions []*gollm.FunctionDefinition) error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "SetFunctionDefinitions", functionDefinitions) ret0, _ := ret[0].(error) return ret0 } // SetFunctionDefinitions indicates an expected call of SetFunctionDefinitions. func (mr *MockChatMockRecorder) SetFunctionDefinitions(functionDefinitions any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetFunctionDefinitions", reflect.TypeOf((*MockChat)(nil).SetFunctionDefinitions), functionDefinitions) } ================================================ FILE: internal/mocks/tools_mock.go ================================================ // Code generated by MockGen. DO NOT EDIT. // Source: github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools (interfaces: Tool) // // Generated by this command: // // mockgen -destination=tools_mock.go -package=mocks github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools Tool // // Package mocks is a generated GoMock package. package mocks import ( context "context" reflect "reflect" gollm "github.com/GoogleCloudPlatform/kubectl-ai/gollm" gomock "go.uber.org/mock/gomock" ) // MockTool is a mock of Tool interface. type MockTool struct { ctrl *gomock.Controller recorder *MockToolMockRecorder isgomock struct{} } // MockToolMockRecorder is the mock recorder for MockTool. type MockToolMockRecorder struct { mock *MockTool } // NewMockTool creates a new mock instance. func NewMockTool(ctrl *gomock.Controller) *MockTool { mock := &MockTool{ctrl: ctrl} mock.recorder = &MockToolMockRecorder{mock} return mock } // EXPECT returns an object that allows the caller to indicate expected use. func (m *MockTool) EXPECT() *MockToolMockRecorder { return m.recorder } // CheckModifiesResource mocks base method. func (m *MockTool) CheckModifiesResource(args map[string]any) string { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "CheckModifiesResource", args) ret0, _ := ret[0].(string) return ret0 } // CheckModifiesResource indicates an expected call of CheckModifiesResource. func (mr *MockToolMockRecorder) CheckModifiesResource(args any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckModifiesResource", reflect.TypeOf((*MockTool)(nil).CheckModifiesResource), args) } // Description mocks base method. func (m *MockTool) Description() string { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Description") ret0, _ := ret[0].(string) return ret0 } // Description indicates an expected call of Description. func (mr *MockToolMockRecorder) Description() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Description", reflect.TypeOf((*MockTool)(nil).Description)) } // FunctionDefinition mocks base method. func (m *MockTool) FunctionDefinition() *gollm.FunctionDefinition { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "FunctionDefinition") ret0, _ := ret[0].(*gollm.FunctionDefinition) return ret0 } // FunctionDefinition indicates an expected call of FunctionDefinition. func (mr *MockToolMockRecorder) FunctionDefinition() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FunctionDefinition", reflect.TypeOf((*MockTool)(nil).FunctionDefinition)) } // IsInteractive mocks base method. func (m *MockTool) IsInteractive(args map[string]any) (bool, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "IsInteractive", args) ret0, _ := ret[0].(bool) ret1, _ := ret[1].(error) return ret0, ret1 } // IsInteractive indicates an expected call of IsInteractive. func (mr *MockToolMockRecorder) IsInteractive(args any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsInteractive", reflect.TypeOf((*MockTool)(nil).IsInteractive), args) } // Name mocks base method. func (m *MockTool) Name() string { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Name") ret0, _ := ret[0].(string) return ret0 } // Name indicates an expected call of Name. func (mr *MockToolMockRecorder) Name() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Name", reflect.TypeOf((*MockTool)(nil).Name)) } // Run mocks base method. func (m *MockTool) Run(ctx context.Context, args map[string]any) (any, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Run", ctx, args) ret0, _ := ret[0].(any) ret1, _ := ret[1].(error) return ret0, ret1 } // Run indicates an expected call of Run. func (mr *MockToolMockRecorder) Run(ctx, args any) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Run", reflect.TypeOf((*MockTool)(nil).Run), ctx, args) } ================================================ FILE: k8s/all_in_one.yaml ================================================ apiVersion: v1 kind: Namespace metadata: name: computer labels: name: computer --- apiVersion: v1 kind: ServiceAccount metadata: name: normal-user namespace: computer --- apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: namespace: computer name: reader-all-but-secrets rules: - apiGroups: [""] resources: ["pods", "pods/log", "pods/status", "configmaps", "persistentvolumeclaims", "replicationcontrollers", "resourcequotas", "limitranges", "endpoints", "events", "services"] verbs: ["get", "list", "watch"] - apiGroups: ["apps"] resources: ["deployments", "daemonsets", "replicasets", "statefulsets"] verbs: ["get", "list", "watch"] - apiGroups: ["autoscaling"] resources: ["horizontalpodautoscalers"] verbs: ["get", "list", "watch"] - apiGroups: ["batch"] resources: ["jobs", "cronjobs"] verbs: ["get", "list", "watch"] - apiGroups: ["extensions"] resources: ["deployments", "daemonsets", "replicasets", "ingresses"] verbs: ["get", "list", "watch"] - apiGroups: ["policy"] resources: ["poddisruptionbudgets"] verbs: ["get", "list", "watch"] - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies", "ingresses"] verbs: ["get", "list", "watch"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: normal-user-reader-binding namespace: computer subjects: - kind: ServiceAccount name: normal-user namespace: computer roleRef: kind: Role name: reader-all-but-secrets apiGroup: rbac.authorization.k8s.io --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRole metadata: name: reader-cluster-resources rules: - apiGroups: ["apiextensions.k8s.io"] resources: ["customresourcedefinitions"] verbs: ["get", "list", "watch"] - apiGroups: [""] resources: ["nodes", "persistentvolumes", "namespaces"] verbs: ["get", "list", "watch"] - apiGroups: ["*"] resources: ["*"] verbs: ["get", "list", "watch"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding metadata: name: normal-user-cluster-reader-binding subjects: - kind: ServiceAccount name: normal-user namespace: computer roleRef: kind: ClusterRole name: reader-cluster-resources apiGroup: rbac.authorization.k8s.io ================================================ FILE: k8s/kubectl-ai-gke.yaml ================================================ apiVersion: v1 kind: Namespace metadata: name: kubectl-ai --- kind: ServiceAccount apiVersion: v1 metadata: name: kubectl-ai namespace: kubectl-ai --- kind: Deployment apiVersion: apps/v1 metadata: name: kubectl-ai namespace: kubectl-ai spec: replicas: 1 selector: matchLabels: app: kubectl-ai template: metadata: labels: app: kubectl-ai spec: serviceAccountName: kubectl-ai containers: - name: kubectl-ai image: REPLACE_WITH_YOUR_IMAGE # e.g. us-central1-docker.pkg.dev/PROJECT_ID/kubectl-ai/kubectl-ai:latest args: - --ui-type=web - --ui-listen-address=0.0.0.0:8080 - --v=4 - --alsologtostderr - --sandbox=k8s env: - name: GOOGLE_CLOUD_PROJECT value: "PROJECT_ID" - name: GOOGLE_CLOUD_LOCATION value: "global" - name: GEMINI_API_KEY value: "REPLACE_WITH_YOUR_GEMINI_API_KEY" --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: kubectl-ai:view namespace: kubectl-ai roleRef: apiGroup: rbac.authorization.k8s.io kind: ClusterRole name: view subjects: - kind: ServiceAccount name: kubectl-ai --- kind: Service apiVersion: v1 metadata: name: kubectl-ai namespace: kubectl-ai labels: app: kubectl-ai spec: selector: app: kubectl-ai ports: - port: 80 targetPort: 8080 protocol: TCP --- # 1. The ClusterRole that grants read-only access to most resources # This is a cluster-wide role, so it does not have a namespace. apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRole metadata: # This name is shared across the cluster. name: read-only-except-secrets-cluster-role rules: - apiGroups: - "" # core API group resources: # List all core resource types EXCEPT "secrets" - configmaps - endpoints - events - limitranges - namespaces - nodes - persistentvolumeclaims - persistentvolumes - pods - podtemplates - replicationcontrollers - resourcequotas - serviceaccounts - services verbs: - get - list - watch - apiGroups: - "*" # All other current and future API groups resources: - "*" # All current and future resources in those groups (including CRDs and CRs) verbs: - get - list - watch --- # 2. The ClusterRoleBinding that connects the ServiceAccount to the Role apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding metadata: name: read-only-kubectl-ai-binding subjects: # Grant the permissions to the specific ServiceAccount in the specific namespace - kind: ServiceAccount name: kubectl-ai namespace: kubectl-ai roleRef: # This refers to the ClusterRole defined above kind: ClusterRole name: read-only-except-secrets-cluster-role apiGroup: rbac.authorization.k8s.io --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRole metadata: name: kubectl-ai-computer-manager rules: - apiGroups: - "" resources: - pods - pods/exec - configmaps - secrets verbs: - create - get - delete --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: kubectl-ai-computer-manager-binding namespace: computer roleRef: apiGroup: rbac.authorization.k8s.io kind: ClusterRole name: kubectl-ai-computer-manager subjects: - kind: ServiceAccount name: kubectl-ai namespace: kubectl-ai ================================================ FILE: k8s/kubectl-ai.yaml ================================================ kind: Deployment apiVersion: apps/v1 metadata: name: kubectl-ai spec: replicas: 1 selector: matchLabels: app: kubectl-ai template: metadata: labels: app: kubectl-ai spec: serviceAccountName: kubectl-ai containers: - name: kubectl-ai image: kubectl-ai:latest args: - --ui-type=web envFrom: - secretRef: name: kubectl-ai --- kind: Secret apiVersion: v1 metadata: name: kubectl-ai labels: app: kubectl-ai type: Opaque --- kind: ServiceAccount apiVersion: v1 metadata: name: kubectl-ai --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: kubectl-ai:view roleRef: apiGroup: rbac.authorization.k8s.io kind: ClusterRole name: view subjects: - kind: ServiceAccount name: kubectl-ai --- kind: Service apiVersion: v1 metadata: name: kubectl-ai labels: app: kubectl-ai spec: selector: app: kubectl-ai ports: - port: 80 targetPort: 8888 protocol: TCP ================================================ FILE: k8s/sandbox/all-in-one.yaml ================================================ apiVersion: v1 kind: Namespace metadata: name: computer labels: name: computer --- apiVersion: v1 kind: ServiceAccount metadata: name: normal-user namespace: computer --- apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: namespace: computer name: reader-all-but-secrets rules: - apiGroups: [""] resources: ["pods", "pods/log", "pods/status", "configmaps", "persistentvolumeclaims", "replicationcontrollers", "resourcequotas", "limitranges", "endpoints", "events", "services"] verbs: ["get", "list", "watch"] - apiGroups: ["apps"] resources: ["deployments", "daemonsets", "replicasets", "statefulsets"] verbs: ["get", "list", "watch"] - apiGroups: ["autoscaling"] resources: ["horizontalpodautoscalers"] verbs: ["get", "list", "watch"] - apiGroups: ["batch"] resources: ["jobs", "cronjobs"] verbs: ["get", "list", "watch"] - apiGroups: ["extensions"] resources: ["deployments", "daemonsets", "replicasets", "ingresses"] verbs: ["get", "list", "watch"] - apiGroups: ["policy"] resources: ["poddisruptionbudgets"] verbs: ["get", "list", "watch"] - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies", "ingresses"] verbs: ["get", "list", "watch"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: normal-user-reader-binding namespace: computer subjects: - kind: ServiceAccount name: normal-user namespace: computer roleRef: kind: Role name: reader-all-but-secrets apiGroup: rbac.authorization.k8s.io --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRole metadata: name: reader-cluster-resources rules: - apiGroups: ["apiextensions.k8s.io"] resources: ["customresourcedefinitions"] verbs: ["get", "list", "watch"] - apiGroups: [""] resources: ["nodes", "persistentvolumes", "namespaces", "pods", "services", "endpoints", "events", "configmaps", "serviceaccounts"] verbs: ["get", "list", "watch"] - apiGroups: ["apps"] resources: ["deployments", "daemonsets", "statefulsets", "replicasets"] verbs: ["get", "list", "watch"] - apiGroups: ["batch"] resources: ["jobs", "cronjobs"] verbs: ["get", "list", "watch"] - apiGroups: ["networking.k8s.io"] resources: ["ingresses", "networkpolicies"] verbs: ["get", "list", "watch"] - apiGroups: ["storage.k8s.io"] resources: ["storageclasses"] verbs: ["get", "list", "watch"] - apiGroups: ["rbac.authorization.k8s.io"] resources: ["roles", "rolebindings", "clusterroles", "clusterrolebindings"] verbs: ["get", "list", "watch"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding metadata: name: normal-user-cluster-reader-binding subjects: - kind: ServiceAccount name: normal-user namespace: computer roleRef: kind: ClusterRole name: reader-cluster-resources apiGroup: rbac.authorization.k8s.io ================================================ FILE: k8s/sandbox/cluster_role.yaml ================================================ apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRole metadata: name: reader-cluster-resources rules: - apiGroups: ["apiextensions.k8s.io"] resources: ["customresourcedefinitions"] verbs: ["get", "list", "watch"] - apiGroups: [""] resources: ["nodes", "persistentvolumes", "namespaces", "pods", "services", "endpoints", "events", "configmaps", "serviceaccounts"] verbs: ["get", "list", "watch"] - apiGroups: ["apps"] resources: ["deployments", "daemonsets", "statefulsets", "replicasets"] verbs: ["get", "list", "watch"] - apiGroups: ["batch"] resources: ["jobs", "cronjobs"] verbs: ["get", "list", "watch"] - apiGroups: ["networking.k8s.io"] resources: ["ingresses", "networkpolicies"] verbs: ["get", "list", "watch"] - apiGroups: ["storage.k8s.io"] resources: ["storageclasses"] verbs: ["get", "list", "watch"] - apiGroups: ["rbac.authorization.k8s.io"] resources: ["roles", "rolebindings", "clusterroles", "clusterrolebindings"] verbs: ["get", "list", "watch"] ================================================ FILE: k8s/sandbox/cluster_role_binding.yaml ================================================ apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRoleBinding metadata: name: normal-user-cluster-reader-binding subjects: - kind: ServiceAccount name: normal-user namespace: computer roleRef: kind: ClusterRole name: reader-cluster-resources apiGroup: rbac.authorization.k8s.io ================================================ FILE: k8s/sandbox/namespace.yaml ================================================ apiVersion: v1 kind: Namespace metadata: name: computer labels: name: computer ================================================ FILE: k8s/sandbox/role.yaml ================================================ apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: namespace: computer name: reader-all-but-secrets rules: - apiGroups: [""] resources: ["pods", "pods/log", "pods/status", "configmaps", "persistentvolumeclaims", "replicationcontrollers", "resourcequotas", "limitranges", "endpoints", "events", "services"] verbs: ["get", "list", "watch"] - apiGroups: ["apps"] resources: ["deployments", "daemonsets", "replicasets", "statefulsets"] verbs: ["get", "list", "watch"] - apiGroups: ["autoscaling"] resources: ["horizontalpodautoscalers"] verbs: ["get", "list", "watch"] - apiGroups: ["batch"] resources: ["jobs", "cronjobs"] verbs: ["get", "list", "watch"] - apiGroups: ["extensions"] resources: ["deployments", "daemonsets", "replicasets", "ingresses"] verbs: ["get", "list", "watch"] - apiGroups: ["policy"] resources: ["poddisruptionbudgets"] verbs: ["get", "list", "watch"] - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies", "ingresses"] verbs: ["get", "list", "watch"] ================================================ FILE: k8s/sandbox/role_binding.yaml ================================================ apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: normal-user-reader-binding namespace: computer subjects: - kind: ServiceAccount name: normal-user namespace: computer roleRef: kind: Role name: reader-all-but-secrets apiGroup: rbac.authorization.k8s.io ================================================ FILE: k8s/sandbox/service_account.yaml ================================================ apiVersion: v1 kind: ServiceAccount metadata: name: normal-user namespace: computer ================================================ FILE: kubectl-utils/README.md ================================================ kubectl-utils contains some experimental kubectl extensions that should help us write simpler evals for kubectl-ai They may one day be useful in their own right, but that is not the current goal. # kubectl expect kubectl expect polls an object, waiting for a CEL expression to be true. Example usage: `kubectl expect StatefulSet/mysql 'self.status.replicas >= 1'` ================================================ FILE: kubectl-utils/cmd/kubectl-expect/main.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package main import ( "context" "flag" "fmt" "os" "strings" "time" "github.com/GoogleCloudPlatform/kubectl-ai/kubectl-utils/pkg/kel" "github.com/GoogleCloudPlatform/kubectl-ai/kubectl-utils/pkg/kube" celtypes "github.com/google/cel-go/common/types" "github.com/spf13/pflag" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" "k8s.io/klog/v2" ) func main() { ctx := context.Background() if err := run(ctx); err != nil { fmt.Fprintf(os.Stderr, "%v\n", err) } } func run(ctx context.Context) error { // log := klog.FromContext(ctx) namespace := "" kubeconfig := "" pflag.StringVarP(&namespace, "namespace", "n", namespace, "If present, the namespace scope for this CLI request") pflag.StringVar(&kubeconfig, "kubeconfig", kubeconfig, "Path to the kubeconfig file to use for CLI requests.") klog.InitFlags(nil) pflag.CommandLine.AddGoFlagSet(flag.CommandLine) pflag.Parse() args := pflag.Args() if len(args) < 2 { return fmt.Errorf("expected [target] [cel-expression]") } target := args[0] celExpressionText := args[1] kubeClient, err := kube.NewClient(kubeconfig) if err != nil { return err } tokens := strings.Split(target, "/") if len(tokens) != 2 { return fmt.Errorf("expected target like Pod/") } // Find the resource (kind) the user is asking about resource, err := kubeClient.FindResource(ctx, tokens[0]) if err != nil { return err } // Compute namespace, defaulting to kubeconfig or default if namespace == "" && resource.Namespaced { namespace, err = kubeClient.DefaultNamespace() if err != nil { return err } } // Compile the CEL expression env, err := kel.NewEnv() if err != nil { return fmt.Errorf("initializing CEL: %w", err) } celExpression, err := kel.NewExpression(env, celExpressionText) if err != nil { return err } // build a pretty-printer for outputting status while polling printer, err := celExpression.BuildStatusPrinter(ctx) if err != nil { return fmt.Errorf("building status printer: %w", err) } // Get ready to get the object id := types.NamespacedName{ Namespace: namespace, Name: tokens[1], } gv := schema.GroupVersion{ Group: resource.Group, Version: resource.Version, } gvr := gv.WithResource(resource.Name) gvk := gv.WithKind(resource.Kind) client := kubeClient.ForGVR(gvr, id.Namespace) // Poll the object until the CEL expression returns true for { // We _could_ watch... time.Sleep(1 * time.Second) u, err := client.Get(ctx, id.Name, metav1.GetOptions{}) if err != nil { return fmt.Errorf("getting %s %s: %w", gvk.Kind, id.Name, err) } out, err := celExpression.Eval(ctx, u) if err != nil { return err } done := false switch out.Type() { case celtypes.BoolType: v := out.Value().(bool) if v { done = true } default: return fmt.Errorf("unhandled type for CEL expression: %v", out.Type()) } if done { break } // Pretty print some intermediate values if we can if printer != nil { s := printer(ctx, u) fmt.Printf("waiting for %q (%s)\n", celExpression.CELText, s) } else { fmt.Printf("waiting for %q\n", celExpression.CELText) } } return nil } ================================================ FILE: kubectl-utils/go.mod ================================================ module github.com/GoogleCloudPlatform/kubectl-ai/kubectl-utils go 1.24.0 toolchain go1.24.3 require ( github.com/google/cel-go v0.25.0 github.com/spf13/pflag v1.0.6 google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2 google.golang.org/protobuf v1.36.6 k8s.io/apimachinery v0.33.0 k8s.io/client-go v0.33.0 k8s.io/klog/v2 v2.130.1 ) require ( cel.dev/expr v0.23.1 // indirect github.com/antlr4-go/antlr/v4 v4.13.0 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/emicklei/go-restful/v3 v3.11.0 // indirect github.com/fxamacker/cbor/v2 v2.7.0 // indirect github.com/go-logr/logr v1.4.2 // indirect github.com/go-openapi/jsonpointer v0.21.0 // indirect github.com/go-openapi/jsonreference v0.20.2 // indirect github.com/go-openapi/swag v0.23.0 // indirect github.com/gogo/protobuf v1.3.2 // indirect github.com/google/gnostic-models v0.6.9 // indirect github.com/google/go-cmp v0.7.0 // indirect github.com/google/uuid v1.6.0 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/stoewer/go-strcase v1.2.0 // indirect github.com/x448/float16 v0.8.4 // indirect golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc // indirect golang.org/x/net v0.38.0 // indirect golang.org/x/oauth2 v0.27.0 // indirect golang.org/x/sys v0.31.0 // indirect golang.org/x/term v0.30.0 // indirect golang.org/x/text v0.23.0 // indirect golang.org/x/time v0.9.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20250428153025-10db94c68c34 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect k8s.io/api v0.33.0 // indirect k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff // indirect k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738 // indirect sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 // indirect sigs.k8s.io/randfill v1.0.0 // indirect sigs.k8s.io/structured-merge-diff/v4 v4.6.0 // indirect sigs.k8s.io/yaml v1.4.0 // indirect ) ================================================ FILE: kubectl-utils/go.sum ================================================ cel.dev/expr v0.23.1 h1:K4KOtPCJQjVggkARsjG9RWXP6O4R73aHeJMa/dmCQQg= cel.dev/expr v0.23.1/go.mod h1:hLPLo1W4QUmuYdA72RBX06QTs6MXw941piREPl3Yfiw= github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI= github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/emicklei/go-restful/v3 v3.11.0 h1:rAQeMHw1c7zTmncogyy8VvRZwtkmkZ4FxERmMY4rD+g= github.com/emicklei/go-restful/v3 v3.11.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= github.com/fxamacker/cbor/v2 v2.7.0 h1:iM5WgngdRBanHcxugY4JySA0nk1wZorNOpTgCMedv5E= github.com/fxamacker/cbor/v2 v2.7.0/go.mod h1:pxXPTn3joSm21Gbwsv0w9OSA2y1HFR9qXEeXQVeNoDQ= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= github.com/go-openapi/jsonpointer v0.21.0 h1:YgdVicSA9vH5RiHs9TZW5oyafXZFc6+2Vc1rr/O9oNQ= github.com/go-openapi/jsonpointer v0.21.0/go.mod h1:IUyH9l/+uyhIYQ/PXVA41Rexl+kOkAPDdXEYns6fzUY= github.com/go-openapi/jsonreference v0.20.2 h1:3sVjiK66+uXK/6oQ8xgcRKcFgQ5KXa2KvnJRumpMGbE= github.com/go-openapi/jsonreference v0.20.2/go.mod h1:Bl1zwGIM8/wsvqjsOQLJ/SH+En5Ap4rVB5KVcIDZG2k= github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+GrE= github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ= github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI= github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= github.com/google/cel-go v0.25.0 h1:jsFw9Fhn+3y2kBbltZR4VEz5xKkcIFRPDnuEzAGv5GY= github.com/google/cel-go v0.25.0/go.mod h1:hjEb6r5SuOSlhCHmFoLzu8HGCERvIsDAbxDAyNU/MmI= github.com/google/gnostic-models v0.6.9 h1:MU/8wDLif2qCXZmzncUQ/BOfxWfthHi63KqpoNbWqVw= github.com/google/gnostic-models v0.6.9/go.mod h1:CiWsm0s6BSQd1hRn8/QmxqB6BesYcbSZxsz9b0KuDBw= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20241029153458-d1b30febd7db h1:097atOisP2aRj7vFgYQBbFN4U4JNXUNYpxael3UzMyo= github.com/google/pprof v0.0.0-20241029153458-d1b30febd7db/go.mod h1:vavhavw2zAxS5dIdcRluK6cSGGPlZynqzFM8NdvU144= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/onsi/ginkgo/v2 v2.21.0 h1:7rg/4f3rB88pb5obDgNZrNHrQ4e6WpjonchcpuBRnZM= github.com/onsi/ginkgo/v2 v2.21.0/go.mod h1:7Du3c42kxCUegi0IImZ1wUQzMBVecgIHjR1C+NkhLQo= github.com/onsi/gomega v1.35.1 h1:Cwbd75ZBPxFSuZ6T+rN/WCb/gOc6YgFBXLlZLhC7Ds4= github.com/onsi/gomega v1.35.1/go.mod h1:PvZbdDc8J6XJEpDK4HCuRBm8a6Fzp9/DmhC9C7yFlog= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o= github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stoewer/go-strcase v1.2.0 h1:Z2iHWqGXH00XYgqDmNgQbIBxf3wrNq0F3feEy0ainaU= github.com/stoewer/go-strcase v1.2.0/go.mod h1:IBiWB2sKIp3wVVQ3Y035++gc+knqhUQag1KpM8ahLw8= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc h1:mCRnTeVUjcrhlRmO0VK8a6k6Rrf6TF9htwo2pJVSjIU= golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc/go.mod h1:V1LtkGg67GoY2N1AnLN78QLrzxkLyJw7RJb1gzOOz9w= golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8= golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= golang.org/x/oauth2 v0.27.0 h1:da9Vo7/tDv5RH/7nZDz1eMGS/q1Vv1N/7FCrBhI9I3M= golang.org/x/oauth2 v0.27.0/go.mod h1:onh5ek6nERTohokkhCD/y2cV4Do3fxFHFuAejCkRWT8= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/term v0.30.0 h1:PQ39fJZ+mfadBm0y5WlL4vlM7Sx1Hgf13sMIY2+QS9Y= golang.org/x/term v0.30.0/go.mod h1:NYYFdzHoI5wRh/h5tDMdMqCqPJZEuNqVR5xJLd/n67g= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY= golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4= golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY= golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.26.0 h1:v/60pFQmzmT9ExmjDv2gGIfi3OqfKoEP6I5+umXlbnQ= golang.org/x/tools v0.26.0/go.mod h1:TPVVj70c7JJ3WCazhD8OdXcZg/og+b9+tH/KxylGwH0= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2 h1:vPV0tzlsK6EzEDHNNH5sa7Hs9bd7iXR7B1tSiPepkV0= google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2/go.mod h1:pKLAc5OolXC3ViWGI62vvC0n10CpwAtRcTNCFwTKBEw= google.golang.org/genproto/googleapis/rpc v0.0.0-20250428153025-10db94c68c34 h1:h6p3mQqrmT1XkHVTfzLdNz1u7IhINeZkz67/xTbOuWs= google.golang.org/genproto/googleapis/rpc v0.0.0-20250428153025-10db94c68c34/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A= google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY= google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/evanphx/json-patch.v4 v4.12.0 h1:n6jtcsulIzXPJaxegRbvFNNrZDjbij7ny3gmSPG+6V4= gopkg.in/evanphx/json-patch.v4 v4.12.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M= gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc= gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/api v0.33.0 h1:yTgZVn1XEe6opVpP1FylmNrIFWuDqe2H0V8CT5gxfIU= k8s.io/api v0.33.0/go.mod h1:CTO61ECK/KU7haa3qq8sarQ0biLq2ju405IZAd9zsiM= k8s.io/apimachinery v0.33.0 h1:1a6kHrJxb2hs4t8EE5wuR/WxKDwGN1FKH3JvDtA0CIQ= k8s.io/apimachinery v0.33.0/go.mod h1:BHW0YOu7n22fFv/JkYOEfkUYNRN0fj0BlvMFWA7b+SM= k8s.io/client-go v0.33.0 h1:UASR0sAYVUzs2kYuKn/ZakZlcs2bEHaizrrHUZg0G98= k8s.io/client-go v0.33.0/go.mod h1:kGkd+l/gNGg8GYWAPr0xF1rRKvVWvzh9vmZAMXtaKOg= k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk= k8s.io/klog/v2 v2.130.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE= k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff h1:/usPimJzUKKu+m+TE36gUyGcf03XZEP0ZIKgKj35LS4= k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff/go.mod h1:5jIi+8yX4RIb8wk3XwBo5Pq2ccx4FP10ohkbSKCZoK8= k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738 h1:M3sRQVHv7vB20Xc2ybTt7ODCeFj6JSWYFzOFnYeS6Ro= k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0= sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 h1:/Rv+M11QRah1itp8VhT6HoVx1Ray9eB4DBr+K+/sCJ8= sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3/go.mod h1:18nIHnGi6636UCz6m8i4DhaJ65T6EruyzmoQqI2BVDo= sigs.k8s.io/randfill v0.0.0-20250304075658-069ef1bbf016/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= sigs.k8s.io/randfill v1.0.0 h1:JfjMILfT8A6RbawdsK2JXGBR5AQVfd+9TbzrlneTyrU= sigs.k8s.io/randfill v1.0.0/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= sigs.k8s.io/structured-merge-diff/v4 v4.6.0 h1:IUA9nvMmnKWcj5jl84xn+T5MnlZKThmUW1TdblaLVAc= sigs.k8s.io/structured-merge-diff/v4 v4.6.0/go.mod h1:dDy58f92j70zLsuZVuUX5Wp9vtxXpaZnkPGWeqDfCps= sigs.k8s.io/yaml v1.4.0 h1:Mk1wCc2gy/F0THH0TAp1QYyJNzRm2KCLy3o5ASXVI5E= sigs.k8s.io/yaml v1.4.0/go.mod h1:Ejl7/uTz7PSA4eKMyQCUTnhZYNmLIl+5c2lQPGR2BPY= ================================================ FILE: kubectl-utils/pkg/kel/expression.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package kel import ( "context" "fmt" "github.com/google/cel-go/cel" celtypes "github.com/google/cel-go/common/types" "github.com/google/cel-go/common/types/ref" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/klog/v2" ) func NewEnv() (*cel.Env, error) { // TODO: Can we / should we do better than AnyType? env, err := cel.NewEnv( cel.Variable("self", cel.AnyType), ) return env, err } type Expression struct { CELText string Program cel.Program AST *cel.Ast Env *cel.Env } func NewExpression(env *cel.Env, celExpression string) (*Expression, error) { ast, issues := env.Compile(celExpression) if issues != nil && issues.Err() != nil { return nil, fmt.Errorf("invalid expression %q: %w", celExpression, issues.Err()) } prg, err := env.Program(ast) if err != nil { return nil, fmt.Errorf("invalid expression %q: %w", celExpression, err) } return &Expression{ CELText: celExpression, AST: ast, Program: prg, Env: env, }, nil } func (x *Expression) Eval(ctx context.Context, self *unstructured.Unstructured) (ref.Val, error) { log := klog.FromContext(ctx) inputs := x.buildInputs(self) out, details, err := x.Program.Eval(inputs) if err != nil { return nil, fmt.Errorf("evaluating CEL expression: %w", err) } log.V(2).Info("evaluated CEL expression", "out", out, "details", details) return out, nil } func (x *Expression) buildInputs(self *unstructured.Unstructured) map[string]any { inputs := map[string]any{ "self": celtypes.NewDynamicMap(&unstructuredToCELAdapter{}, self.Object), } return inputs } type unstructuredToCELAdapter struct { } func (a *unstructuredToCELAdapter) NativeToValue(value any) ref.Val { switch value := value.(type) { case string: return celtypes.String(value) case int: return celtypes.Int(value) case int64: return celtypes.Int(value) case map[string]any: return celtypes.NewDynamicMap(a, value) default: klog.Fatalf("unhandled type %T", value) return nil } } ================================================ FILE: kubectl-utils/pkg/kel/info.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package kel import ( "context" "fmt" "strings" "github.com/google/cel-go/cel" exprpb "google.golang.org/genproto/googleapis/api/expr/v1alpha1" "google.golang.org/protobuf/proto" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/klog/v2" ) type InfoFunction func(ctx context.Context, self *unstructured.Unstructured) string // BuildStatusPrinter returns an InfoFunction that attempts to report important values from the evaluation of the CEL expression func (x *Expression) BuildStatusPrinter(ctx context.Context) (InfoFunction, error) { log := klog.FromContext(ctx) checkedExpr, err := cel.AstToCheckedExpr(x.AST) if err != nil { return nil, fmt.Errorf("parsing CEL ast: %w", err) } v := checkedExpr.Expr.ExprKind switch v := v.(type) { case *exprpb.Expr_CallExpr: printFunction := "" switch v.CallExpr.Function { case "_==_": printFunction = "=" case "_>=_": printFunction = ">=" case "_<=_": printFunction = "<=" case "_>_": printFunction = ">" case "_<_": printFunction = "<" default: klog.Warningf("unhandled function %q", v.CallExpr.Function) return nil, nil } log.V(2).Info("recognized function", "function", printFunction) return x.buildFunctionPrinterFor(v.CallExpr.Args) default: klog.Warningf("unhandled expression kind %T", checkedExpr.Expr.ExprKind) return nil, nil } } func (x *Expression) buildFunctionPrinterFor(args []*exprpb.Expr) (InfoFunction, error) { checkedExpr, err := cel.AstToCheckedExpr(x.AST) if err != nil { return nil, fmt.Errorf("parsing CEL ast: %w", err) } type debugValue struct { Key string Program cel.Program } var debugValues []debugValue for _, arg := range args { shouldPrint := true v := arg.ExprKind switch v := v.(type) { case *exprpb.Expr_ConstExpr: // Don't print constants, 2=2 is not informative shouldPrint = false case *exprpb.Expr_SelectExpr: shouldPrint = true default: klog.Warningf("unhandled expression kind %T", v) } if !shouldPrint { continue } checkedArg := proto.Clone(checkedExpr).(*exprpb.CheckedExpr) checkedArg.Expr = arg ast := cel.CheckedExprToAst(checkedArg) celExpression, err := cel.AstToString(ast) if err != nil { return nil, fmt.Errorf("converting expression to string: %w", err) } compiled, issues := x.Env.Compile(celExpression) if issues != nil && issues.Err() != nil { return nil, fmt.Errorf("invalid expression %q: %w", celExpression, issues.Err()) } prg, err := x.Env.Program(compiled) if err != nil { return nil, fmt.Errorf("invalid expression %q: %w", celExpression, err) } debugValues = append(debugValues, debugValue{ Key: celExpression, Program: prg, }) } if len(debugValues) == 0 { return nil, nil } return func(ctx context.Context, self *unstructured.Unstructured) string { log := klog.FromContext(ctx) inputs := x.buildInputs(self) var values []string for _, debugValue := range debugValues { s := "" out, details, err := debugValue.Program.Eval(inputs) log.V(2).Info("evaluated CEL expression", "out", out, "details", details, "error", err) if err == nil { s = fmt.Sprintf("%s=%v", debugValue.Key, out.Value()) } else { s = fmt.Sprintf("%s=%v", debugValue.Key, "???") } values = append(values, s) } return strings.Join(values, "; ") }, nil } ================================================ FILE: kubectl-utils/pkg/kube/client.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package kube import ( "fmt" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/client-go/discovery" "k8s.io/client-go/dynamic" "k8s.io/client-go/rest" "k8s.io/client-go/tools/clientcmd" ) // Client is a facade around the various kube interfaces type Client struct { clientConfig clientcmd.ClientConfig DyanmicClient dynamic.Interface DiscoveryClient discovery.DiscoveryInterface } func NewClient(kubeconfig string) (*Client, error) { clientConfig, err := loadKubeconfig(kubeconfig) if err != nil { return nil, err } restConfig, err := clientConfig.ClientConfig() if err != nil { return nil, fmt.Errorf("building kubernetes API configuration: %w", err) } httpClient, err := rest.HTTPClientFor(restConfig) if err != nil { return nil, fmt.Errorf("building http client for rest config: %w", err) } dynamicClient, err := dynamic.NewForConfigAndClient(restConfig, httpClient) if err != nil { return nil, fmt.Errorf("building dynamic client: %w", err) } discoveryClient, err := buildDiscoveryClient(restConfig, httpClient) if err != nil { return nil, err } return &Client{ clientConfig: clientConfig, DyanmicClient: dynamicClient, DiscoveryClient: discoveryClient, }, nil } func loadKubeconfig(kubeconfigPath string) (clientcmd.ClientConfig, error) { rules := clientcmd.NewDefaultClientConfigLoadingRules() if kubeconfigPath != "" { rules.ExplicitPath = kubeconfigPath } clientConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( rules, &clientcmd.ConfigOverrides{}, ) return clientConfig, nil } func (c *Client) DefaultNamespace() (string, error) { ns, _, err := c.clientConfig.Namespace() if err != nil { return "", fmt.Errorf("getting namespace from kubeconfig: %w", err) } namespace := ns if namespace == "" { namespace = "default" } return namespace, nil } // ForGVR returns a dynamic client for the specified GroupVersionResource and namespace func (c *Client) ForGVR(gvr schema.GroupVersionResource, namespace string) dynamic.ResourceInterface { var client dynamic.ResourceInterface if namespace != "" { client = c.DyanmicClient.Resource(gvr).Namespace(namespace) } else { client = c.DyanmicClient.Resource(gvr) } return client } ================================================ FILE: kubectl-utils/pkg/kube/discovery.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package kube import ( "context" "fmt" "net/http" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/client-go/discovery" "k8s.io/client-go/rest" ) func buildDiscoveryClient(restConfig *rest.Config, httpClient *http.Client) (discovery.DiscoveryInterface, error) { // TODO: share cache with kubectl? client, err := discovery.NewDiscoveryClientForConfigAndClient(restConfig, httpClient) if err != nil { return nil, fmt.Errorf("building discovery client: %w", err) } return client, nil } func (c *Client) FindResource(ctx context.Context, name string) (*metav1.APIResource, error) { var matches []metav1.APIResource resourceLists, err := c.DiscoveryClient.ServerPreferredResources() if err != nil { return nil, fmt.Errorf("doing server discovery: %w", err) } for _, resourceList := range resourceLists { gv, err := schema.ParseGroupVersion(resourceList.GroupVersion) if err != nil { return nil, fmt.Errorf("parsing group version %q: %w", resourceList.GroupVersion, err) } for _, resource := range resourceList.APIResources { if resource.Kind == name { if resource.Group == "" { resource.Group = gv.Group } if resource.Version == "" { resource.Version = gv.Version } matches = append(matches, resource) } } } if len(matches) == 0 { return nil, fmt.Errorf("no match for resource %q", name) } if len(matches) > 1 { return nil, fmt.Errorf("found multiple matches for resource %q", name) } resource := matches[0] return &resource, nil } ================================================ FILE: makefile ================================================ # Makefile for kubectl-ai # # This Makefile provides a set of commands to build, test, run, # and manage the kubectl-ai project. # Default target to run when no target is specified. .DEFAULT_GOAL := help # --- Variables --- # Define common variables to avoid repetition and ease maintenance. BIN_DIR := ./bin CMD_DIR := ./cmd BINARY_NAME := kubectl-ai BINARY_PATH := $(BIN_DIR)/$(BINARY_NAME) # Attempt to determine GOPATH/bin for installation. # Fallback to a common default if `go env GOPATH` fails or is empty. GOPATH_BIN := $(shell go env GOPATH)/bin ifeq ($(GOPATH_BIN),/bin) GOPATH_BIN := $(HOME)/go/bin endif # --- Environment Variables from .env --- # If a .env file exists, include it. This makes variables defined in .env # (e.g., API_KEY=123) available as Make variables. # Then, export these variables so they are available in the environment # for shell commands executed by Make recipes. ifneq ($(wildcard .env),) include .env # Extract variable names from .env and export them. # This assumes .env contains lines like VAR=value. ENV_VARS_TO_EXPORT := $(shell awk -F= '{print $$1}' .env | xargs) export $(ENV_VARS_TO_EXPORT) endif # --- Help Target --- # Displays a list of available targets and their descriptions. # Descriptions are extracted from comments following '##'. help: @echo "kubectl-ai Makefile" @echo "-------------------" @echo "Available targets:" @awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z0-9_-]+:.*?## / {printf " %-20s %s\n", $$1, $$2}' $(MAKEFILE_LIST) # --- Build Tasks --- build-recursive: ## Build the binary using dev script (recursive for all modules) @echo "λ Building all modules (recursive using dev script)..." mkdir -p $(BIN_DIR) ./dev/ci/presubmits/go-build.sh build: ## Build single binary for the current platform @echo "λ Building $(BINARY_NAME) for current platform..." mkdir -p $(BIN_DIR) go build -o $(BINARY_PATH) $(CMD_DIR) # --- Run Tasks --- run: ## Run the application @echo "λ Running $(BINARY_NAME) from source..." go run $(CMD_DIR) run-html: ## Run with HTML UI @echo "λ Running $(BINARY_NAME) with HTML UI from source..." go run $(CMD_DIR) --ui-type web # --- Code Quality Tasks (using dev scripts) --- fmt: ## Format code using dev script @echo "λ Formatting code (using dev script)..." ./dev/tasks/format.sh vet: ## Run go vet using dev script @echo "λ Running go vet (using dev script)..." ./dev/ci/presubmits/go-vet.sh tidy: ## Tidy go modules using dev script @echo "λ Tidying go modules (using dev script)..." ./dev/tasks/gomod.sh # --- Verification Tasks (CI-style checks using dev scripts) --- verify-format: ## Verify code formatting @echo "λ Verifying code formatting..." ./dev/ci/presubmits/verify-format.sh verify-gomod: ## Verify go.mod files are tidy @echo "λ Verifying go.mod files..." ./dev/ci/presubmits/verify-gomod.sh verify-autogen: ## Verify auto-generated files are up to date @echo "λ Verifying auto-generated files..." ./dev/ci/presubmits/verify-autogen.sh generate: go generate ./internal/mocks verify-mocks: @echo "λ Verifying mocks..." ./dev/ci/presubmits/verify-mocks.sh # --- Generation Tasks --- generate-actions: ## Generate GitHub Actions workflows @echo "λ Generating GitHub Actions workflows..." ./dev/tasks/generate-github-actions.sh # --- Evaluation Tasks --- run-evals: ## Run evaluations (periodic task) @echo "λ Running evaluations..." ./dev/ci/periodics/run-evals.sh analyze-evals: ## Analyze evaluations (periodic task) @echo "λ Analyzing evaluations..." ./dev/ci/periodics/analyze-evals.sh $(ARGS) # --- Combined Tasks --- # 'check' depends on other verification tasks. They will run as prerequisites. check: verify-format verify-gomod verify-autogen build-recursive vet ## Run all verification checks (presubmit-style) @echo "λ All checks completed." # --- Development Workflow --- # 'dev' and 'dev-html' depend on the 'build' target. dev: build ## Development mode - build and run @echo "λ Starting $(BINARY_NAME) in dev mode..." $(BINARY_PATH) dev-html: build ## Development mode - build and run with HTML UI @echo "λ Starting $(BINARY_NAME) with HTML UI in dev mode..." $(BINARY_PATH) --ui-type web # --- Maintenance Tasks --- clean: ## Clean build artifacts and coverage files @echo "λ Cleaning build artifacts..." rm -rf $(BIN_DIR) rm -f coverage.out coverage.html deps: ## Download Go module dependencies @echo "λ Downloading Go module dependencies..." go mod download update-deps: ## Update Go module dependencies and then tidy @echo "λ Updating Go module dependencies..." go get -u ./... @echo "λ Tidying modules after update..." $(MAKE) tidy # --- Installation --- # 'install' depends on the 'build' target. install: build ## Install the binary to $(GOPATH_BIN) @echo "λ Installing $(BINARY_NAME) to $(GOPATH_BIN)..." cp $(BINARY_PATH) $(GOPATH_BIN)/ @echo "$(BINARY_NAME) installed." # --- Testing --- test: ## Run tests @echo "λ Running tests..." go test ./... test-verbose: ## Run tests with verbose output @echo "λ Running tests (verbose)..." go test -v ./... test-coverage: ## Run tests with coverage and generate HTML report @echo "λ Running tests with coverage..." go test -coverprofile=coverage.out ./... @echo "λ Generating coverage HTML report..." go tool cover -html=coverage.out -o coverage.html @echo "Coverage report generated: coverage.html" ================================================ FILE: modelserving/.gitignore ================================================ .build .cache ================================================ FILE: modelserving/README.md ================================================ # Model Serving This directory provides components to build and deploy Large Language Model (LLM) serving endpoints. - [`k8s/`](k8s/): Kubernetes manifests for model serving components. - [`images/`](images/): Dockerfiles for building model serving container images. - [`dev/tasks`](dev/tasks): Development-related scripts for model serving. - `download-model`: fetch the required model weights (e.g., Gemma 3 12B IT). - `build-images`: runs `download-model`, and then build the Docker image using the provided Dockerfile in `images/`. - `deploy-to-gke` or `dev/tasks/deploy-to-kind`: runs `build-images`, and then deploy the model serving Kubernetes manifests to Google Kubernetes Engine (GKE) or a local KinD cluster. Once deployed, the model server will be accessible via a Kubernetes Service defined in the manifest. You can use `kubectl get svc` to find the service details and access its endpoint. - `run-local`: run the model server locally for testing purposes, bypassing Kubernetes. ================================================ FILE: modelserving/dev/tasks/build-images ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT}/modelserving cd "${SRC_DIR}" if [[ -z "${IMAGE_PREFIX:-}" ]]; then IMAGE_PREFIX="" fi echo "Building images with prefix ${IMAGE_PREFIX}" if [[ -z "${TAG:-}" ]]; then TAG=latest fi if [[ -z "${ARCHITECTURES:-}" ]]; then ARCHITECTURES=cpu,cuda fi echo "Building for architectures: ${ARCHITECTURES}" LLAMACPP_TAG=b4957 echo "Building llama.cpp version ${LLAMACPP_TAG}" function build_for_architecture() { a=",${ARCHITECTURES:-}," if [[ "${a}" =~ ",${1}," ]]; then return 0 fi return 1 } if [[ -z "${BUILDX_ARGS:-}" ]]; then BUILDX_ARGS="--load" fi dev/tasks/download-model # Note we do not push or load the "base" llama-server images (we do not pass BUILDX_ARGS) # This is because this is only an intermediate image (e.g. used for the gemma3-12b-it image) if build_for_architecture cpu; then docker buildx build \ -f images/llamacpp-server/Dockerfile \ --target llamacpp-server \ -t llamacpp-server-cpu:${TAG} \ --build-arg BASE_IMAGE=debian:latest \ --build-arg BUILDER_IMAGE=debian:latest \ --build-arg "CMAKE_ARGS=-DGGML_RPC=ON" \ --progress=plain . fi # We're running distributed now, so the "worker" nodes need CUDA (rpc-server image), the "head" nodes do not (llamacpp-server image). # # -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined allows us to build in a container without all the CUDA libraries present # # These flags mirror the flags in the llama.cpp github-action: https://github.com/ggml-org/llama.cpp/blob/master/.github/workflows/build.yml # docker buildx build \ # -f images/llamacpp-server/Dockerfile \ # --target llamacpp-server \ # -t llamacpp-server-cuda:${TAG} \ # --build-arg BASE_IMAGE=nvidia/cuda:12.6.2-runtime-ubuntu24.04 \ # --build-arg BUILDER_IMAGE=nvidia/cuda:12.6.2-devel-ubuntu24.04 \ # --build-arg "CMAKE_ARGS=-DGGML_RPC=ON -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=all -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined" \ # --progress=plain . # fi # Build a head node that embeds gemma3 if build_for_architecture cpu; then docker buildx build ${BUILDX_ARGS} \ -f images/llamacpp-gemma3-12b-it/Dockerfile \ -t ${IMAGE_PREFIX}llamacpp-gemma3-12b-it-cpu:${TAG} \ --build-arg BASE_IMAGE=llamacpp-server-cpu:${TAG} \ --progress=plain . fi # Build a worker node that runs rpc-server with CPU support if build_for_architecture cpu; then docker buildx build ${BUILDX_ARGS} \ -f images/llamacpp-server/Dockerfile \ --target rpc-server \ -t ${IMAGE_PREFIX}rpc-server-cpu:${TAG} \ --build-arg BASE_IMAGE=debian:latest \ --build-arg BUILDER_IMAGE=debian:latest \ --build-arg "CMAKE_ARGS=-DGGML_RPC=ON" \ --progress=plain . fi # Build a worker node that runs rpc-server with CUDA support if build_for_architecture cuda; then # -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined allows us to build in a container without all the CUDA libraries present # These flags mirror the flags in the llama.cpp github-action: https://github.com/ggml-org/llama.cpp/blob/master/.github/workflows/build.yml docker buildx build ${BUILDX_ARGS} \ -f images/llamacpp-server/Dockerfile \ --target rpc-server \ -t ${IMAGE_PREFIX}rpc-server-cuda:${TAG} \ --build-arg BASE_IMAGE=nvidia/cuda:12.6.2-runtime-ubuntu24.04 \ --build-arg BUILDER_IMAGE=nvidia/cuda:12.6.2-devel-ubuntu24.04 \ --build-arg "CMAKE_ARGS=-DGGML_RPC=ON -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=all -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined" \ --progress=plain . fi ================================================ FILE: modelserving/dev/tasks/deploy-to-gke ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT}/modelserving cd "${SRC_DIR}" if [[ -z "${GCP_PROJECT_ID:-}" ]]; then GCP_PROJECT_ID=$(gcloud config get project) fi echo "Using GCP_PROJECT_ID=${GCP_PROJECT_ID}" if [[ -z "${KUBE_CONTEXT:-}" ]]; then echo "Listing GKE clusters in project ${GCP_PROJECT_ID}:" gcloud container clusters list --project=${GCP_PROJECT_ID} echo "" echo "Please set CONTEXT to kubectl context to use" exit 1 fi # Pick a probably-unique tag export TAG=`date +%Y%m%d%H%M%S` # Build the image echo "Building images" export IMAGE_PREFIX=gcr.io/${GCP_PROJECT_ID}/ ARCHITECTURES=cpu,cuda BUILDX_ARGS=--push dev/tasks/build-images # TODO: support cpu on GKE? MODEL_IMAGE="${IMAGE_PREFIX}llamacpp-gemma3-12b-it-cpu:${TAG}" RPCSERVER_IMAGE="${IMAGE_PREFIX:-}rpc-server-cuda:${TAG}" # Deploy manifests echo "Deploying manifests" cat k8s/llm-server-rpc.yaml | sed s@llamacpp-gemma3-12b-it-cpu:latest@${MODEL_IMAGE}@g | \ kubectl apply --context=${KUBE_CONTEXT} --server-side -f - cat k8s/rpc-server-cuda.yaml | sed s@rpc-server-cuda:latest@${RPCSERVER_IMAGE}@g | \ kubectl apply --context=${KUBE_CONTEXT} --server-side -f - ================================================ FILE: modelserving/dev/tasks/deploy-to-kind ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT}/modelserving cd "${SRC_DIR}" # Pick a probably-unique tag export TAG=`date +%Y%m%d%H%M%S` # If we're building for kind, default to only building on cpu if [[ -z "${ARCHITECTURES:-}" ]]; then ARCHITECTURES=cpu export ARCHITECTURES fi if [[ -z "${KUBE_CONTEXT:-}" ]]; then KUBE_CONTEXT=kind-kind echo "Defaulting to kube context: ${KUBE_CONTEXT}" fi # Build the image echo "Building images" export IMAGE_PREFIX=fake.registry/ BUILDX_ARGS=--load dev/tasks/build-images MODEL_IMAGE="${IMAGE_PREFIX:-}llamacpp-gemma3-12b-it-cpu:${TAG}" RPCSERVER_IMAGE="${IMAGE_PREFIX:-}rpc-server-cpu:${TAG}" # Load the image into kind echo "Loading images into kind: ${MODEL_IMAGE}, ${RPCSERVER_IMAGE}" kind load docker-image ${MODEL_IMAGE} ${RPCSERVER_IMAGE} # Deploy manifests echo "Deploying manifests" cat k8s/llm-server-cpu.yaml | sed s@llamacpp-gemma3-12b-it-cpu:latest@${MODEL_IMAGE}@g | \ kubectl apply --context=${KUBE_CONTEXT} --server-side -f - cat k8s/rpc-server-cpu.yaml | sed s@rpc-server-cpu:latest@${RPCSERVER_IMAGE}@g | \ kubectl apply --context=${KUBE_CONTEXT} --server-side -f - cat k8s/llm-server-rpc.yaml | sed s@llamacpp-gemma3-12b-it-cpu:latest@${MODEL_IMAGE}@g | \ kubectl apply --context=${KUBE_CONTEXT} --server-side -f - ================================================ FILE: modelserving/dev/tasks/download-model ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT}/modelserving cd "${SRC_DIR}" mkdir -p .cache MODEL_NAME="gemma-3-12b-it-Q4_K_M.gguf" MODEL_PATH=".cache/${MODEL_NAME}" MODEL_REPO="unsloth/gemma-3-12b-it-GGUF" MODEL_URL="https://huggingface.co/${MODEL_REPO}/resolve/main/${MODEL_NAME}" # Fetch SHA-256 checksum from Hugging Face API fetch_expected_checksum() { echo "Fetching expected checksum for ${MODEL_NAME} from Hugging Face..." EXPECTED_SHA256=$(curl -s "https://huggingface.co/api/models/${MODEL_REPO}" | \ jq -r ".siblings[] | select(.rfilename == \"${MODEL_NAME}\") | .sha256") if [[ -z "${EXPECTED_SHA256}" || "${EXPECTED_SHA256}" == "null" ]]; then echo "Failed to retrieve expected SHA256 checksum from Hugging Face" exit 1 fi echo "Expected SHA256: ${EXPECTED_SHA256}" } download_model() { echo "Downloading ${MODEL_NAME}..." wget "${MODEL_URL}" -O "${MODEL_PATH}" } verify_checksum() { echo "Verifying checksum for ${MODEL_NAME}..." local actual_hash actual_hash=$(sha256sum "${MODEL_PATH}" | awk '{print $1}') if [[ "${actual_hash}" != "${EXPECTED_SHA256}" ]]; then echo "Checksum mismatch" echo "Expected: ${EXPECTED_SHA256}" echo "Actual: ${actual_hash}" rm -f "${MODEL_PATH}" exit 1 fi echo "Checksum verified" } # Main logic fetch_expected_checksum if [[ ! -f "${MODEL_PATH}" ]]; then download_model verify_checksum else echo "${MODEL_NAME} already exists. Verifying checksum..." verify_checksum fi ================================================ FILE: modelserving/dev/tasks/run-local ================================================ #!/usr/bin/env bash set -o errexit set -o nounset set -o pipefail REPO_ROOT="$(git rev-parse --show-toplevel)" SRC_DIR=${REPO_ROOT}/modelserving cd "${SRC_DIR}" export ARCHITECTURES=cpu # TODO: we could support cuda locally, but it is slow to build.. # Build and export the docker image; so we are consistent in how we build [[ -x dev/tasks/build-images ]] || { echo "ERROR: build-images script not found or not executable" exit 1 } mkdir -p .build/llamacpp-server-cpu BUILDX_ARGS="--output type=local,dest=.build/llamacpp-server-cpu" dev/tasks/build-images # Default model export LLAMA_ARG_MODEL=${SRC_DIR}/.cache/gemma-3-12b-it-Q4_K_M.gguf # Bigger context size (though not too large given memory) export LLAMA_ARG_CTX_SIZE=16384 LD_LIBRARY_PATH=.build/llamacpp-server-cpu/lib/ .build/llamacpp-server-cpu/llama-server --jinja -fa ================================================ FILE: modelserving/images/llamacpp-gemma3-12b-it/Dockerfile ================================================ ARG BASE_IMAGE FROM ${BASE_IMAGE} # TODO: Add checksum COPY .cache/gemma-3-12b-it-Q4_K_M.gguf /gemma-3-12b-it-Q4_K_M.gguf # Default model ENV LLAMA_ARG_MODEL=/gemma-3-12b-it-Q4_K_M.gguf # Bigger context size (though not too large given memory) ENV LLAMA_ARG_CTX_SIZE=16384 ENTRYPOINT [ "/llama-server" ] ================================================ FILE: modelserving/images/llamacpp-server/Dockerfile ================================================ ARG BUILDER_IMAGE ARG BASE_IMAGE FROM ${BUILDER_IMAGE} AS builder ARG CMAKE_ARGS ARG LLAMACPP_TAG RUN apt-get update RUN apt-get install -y g++ git cmake libcurl4-openssl-dev WORKDIR /src RUN git clone https://github.com/ggml-org/llama.cpp WORKDIR /src/llama.cpp RUN git checkout ${LLAMACPP_TAG} RUN cmake -DCMAKE_BUILD_TYPE=Release ${CMAKE_ARGS} . RUN cmake --build . -j16 --config Release --target llama-server --target rpc-server RUN ldd /src/llama.cpp/bin/llama-server FROM ${BASE_IMAGE} AS rpc-server RUN apt-get update && apt-get install --yes libgomp1 COPY --from=builder /src/llama.cpp/bin/rpc-server /rpc-server COPY --from=builder /src/llama.cpp/bin/lib*.so /lib/ ENTRYPOINT [ "/rpc-server" ] FROM ${BASE_IMAGE} AS llamacpp-server RUN apt-get update && apt-get install --yes libgomp1 COPY --from=builder /src/llama.cpp/bin/llama-server /llama-server COPY --from=builder /src/llama.cpp/bin/lib*.so /lib/ ENTRYPOINT [ "/llama-server" ] ================================================ FILE: modelserving/k8s/llm-server-cpu.yaml ================================================ kind: Deployment apiVersion: apps/v1 metadata: name: llm-server spec: replicas: 1 selector: matchLabels: app: llm-server template: metadata: labels: app: llm-server spec: serviceAccountName: llm-server containers: - name: llm-server image: llamacpp-gemma3-12b-it-cpu:latest # placeholder value, replaced by deployment scripts env: - name: LLAMA_ARG_FLASH_ATTN value: "yes" args: - --jinja # Needed for tool use, no env var --- kind: ServiceAccount apiVersion: v1 metadata: name: llm-server --- kind: Service apiVersion: v1 metadata: name: llm-server labels: app: llm-server spec: selector: app: llm-server ports: - port: 80 targetPort: 8080 protocol: TCP ================================================ FILE: modelserving/k8s/llm-server-rpc.yaml ================================================ kind: Deployment apiVersion: apps/v1 metadata: name: llm-server-rpc spec: replicas: 1 selector: matchLabels: app: llm-server-rpc template: metadata: labels: app: llm-server-rpc spec: serviceAccountName: llm-server-rpc containers: - name: llm-server-rpc image: llamacpp-gemma3-12b-it-cpu:latest # placeholder value, replaced by deployment scripts env: - name: LLAMA_ARG_N_GPU_LAYERS value: "99" - name: LLAMA_ARG_FLASH_ATTN value: "yes" # - name: SERVERS # value: rpc-server-0,rpc-server-1,rpc-server-2,rpc-server-3 args: - --jinja # Needed for tool use, no env var - --rpc - rpc-server-0.rpc-server:50052,rpc-server-1.rpc-server:50052,rpc-server-2.rpc-server:50052,rpc-server-3.rpc-server:50052 --- kind: ServiceAccount apiVersion: v1 metadata: name: llm-server-rpc --- kind: Service apiVersion: v1 metadata: name: llm-server-rpc labels: app: llm-server-rpc spec: selector: app: llm-server-rpc ports: - name: http port: 80 targetPort: 8080 protocol: TCP ================================================ FILE: modelserving/k8s/llm-server.yaml ================================================ kind: Deployment apiVersion: apps/v1 metadata: name: llm-server spec: replicas: 1 selector: matchLabels: app: llm-server template: metadata: labels: app: llm-server spec: serviceAccountName: llm-server containers: - name: llm-server image: llamacpp-gemma3-12b-it-cuda:latest # placeholder value, replaced by deployment scripts env: - name: LLAMA_ARG_N_GPU_LAYERS value: "99" - name: LLAMA_ARG_FLASH_ATTN value: "yes" args: - --jinja # Needed for tool use, no env var resources: limits: nvidia.com/gpu: "1" requests: nvidia.com/gpu: "1" nodeSelector: cloud.google.com/gke-accelerator: nvidia-l4 --- kind: ServiceAccount apiVersion: v1 metadata: name: llm-server --- kind: Service apiVersion: v1 metadata: name: llm-server labels: app: llm-server spec: selector: app: llm-server ports: - port: 80 targetPort: 8080 protocol: TCP ================================================ FILE: modelserving/k8s/rpc-server-cpu.yaml ================================================ kind: ServiceAccount apiVersion: v1 metadata: name: rpc-server --- kind: Service apiVersion: v1 metadata: name: rpc-server labels: app: rpc-server spec: clusterIP: None selector: app: rpc-server --- kind: StatefulSet apiVersion: apps/v1 metadata: name: rpc-server spec: podManagementPolicy: "Parallel" replicas: 4 selector: matchLabels: app: rpc-server serviceName: rpc-server template: metadata: labels: app: rpc-server spec: serviceAccountName: rpc-server containers: - name: rpc-server image: rpc-server-cpu:latest # placeholder value, replaced by deployment scripts args: - --host - 0.0.0.0 - --mem - "4192" ================================================ FILE: modelserving/k8s/rpc-server-cuda.yaml ================================================ kind: ServiceAccount apiVersion: v1 metadata: name: rpc-server --- kind: Service apiVersion: v1 metadata: name: rpc-server labels: app: rpc-server spec: clusterIP: None selector: app: rpc-server --- kind: StatefulSet apiVersion: apps/v1 metadata: name: rpc-server spec: podManagementPolicy: "Parallel" replicas: 4 selector: matchLabels: app: rpc-server serviceName: rpc-server template: metadata: labels: app: rpc-server spec: serviceAccountName: rpc-server containers: - name: rpc-server image: rpc-server-cuda:latest # placeholder value, replaced by deployment scripts args: - --host - 0.0.0.0 resources: limits: nvidia.com/gpu: "1" requests: nvidia.com/gpu: "1" nodeSelector: cloud.google.com/gke-accelerator: nvidia-l4 ================================================ FILE: pkg/agent/agent_e2e_test.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package agent import ( "context" "testing" "time" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/internal/mocks" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sessions" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" "go.uber.org/mock/gomock" ) func recvMsg(t *testing.T, ctx context.Context, ch <-chan any) *api.Message { t.Helper() select { case v := <-ch: m, ok := v.(*api.Message) if !ok { t.Fatalf("recvMsg: expected *api.Message, got %T", v) return nil } return m case <-ctx.Done(): t.Fatalf("timed out waiting for message: %v", ctx.Err()) return nil } } func recvUntil(t *testing.T, ctx context.Context, ch <-chan any, pred func(*api.Message) bool) *api.Message { t.Helper() for { select { case v := <-ch: m, ok := v.(*api.Message) if !ok { t.Fatalf("recvUntil: expected *api.Message, got %T", v) return nil } if pred(m) { return m } case <-ctx.Done(): t.Fatalf("timed out waiting for matching message: %v", ctx.Err()) } } } type fakePart struct { text string calls []gollm.FunctionCall } func (p fakePart) AsText() (string, bool) { if p.text != "" { return p.text, true } return "", false } func (p fakePart) AsFunctionCalls() ([]gollm.FunctionCall, bool) { if p.calls != nil { return p.calls, true } return nil, false } type fakeCandidate struct{ parts []gollm.Part } func (c fakeCandidate) String() string { return "" } func (c fakeCandidate) Parts() []gollm.Part { return c.parts } type fakeChatResponse struct{ candidate gollm.Candidate } func (r fakeChatResponse) UsageMetadata() any { return nil } func (r fakeChatResponse) Candidates() []gollm.Candidate { return []gollm.Candidate{r.candidate} } func fCalls(name string, args map[string]any) gollm.Part { return fakePart{calls: []gollm.FunctionCall{{ID: "1", Name: name, Arguments: args}}} } func fText(s string) gollm.Part { return fakePart{text: s} } func chatWith(parts ...gollm.Part) gollm.ChatResponse { return fakeChatResponse{candidate: fakeCandidate{parts: parts}} } func TestAgentEndToEndToolExecution(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() store := sessions.NewInMemoryChatStore() client := mocks.NewMockClient(ctrl) chat := mocks.NewMockChat(ctrl) client.EXPECT().StartChat(gomock.Any(), "test-model").Return(chat) chat.EXPECT().Initialize(gomock.Any()).Return(nil) chat.EXPECT().SetFunctionDefinitions(gomock.Any()).Return(nil) firstResp := chatWith(fCalls("mocktool", map[string]any{"command": "do"})) secondResp := chatWith(fText("all done")) firstIter := gollm.ChatResponseIterator(func(yield func(gollm.ChatResponse, error) bool) { yield(firstResp, nil) }) secondIter := gollm.ChatResponseIterator(func(yield func(gollm.ChatResponse, error) bool) { yield(secondResp, nil) }) gomock.InOrder( chat.EXPECT().SendStreaming(gomock.Any(), gomock.Any()).Return(firstIter, nil), chat.EXPECT().SendStreaming(gomock.Any(), gomock.Any()).Return(secondIter, nil), ) tool := mocks.NewMockTool(ctrl) tool.EXPECT().Name().Return("mocktool").AnyTimes() tool.EXPECT().Description().Return("mock tool").AnyTimes() tool.EXPECT().FunctionDefinition().Return(&gollm.FunctionDefinition{Name: "mocktool"}).AnyTimes() tool.EXPECT().IsInteractive(gomock.Any()).Return(false, nil).AnyTimes() tool.EXPECT().CheckModifiesResource(gomock.Any()).Return("yes").AnyTimes() tool.EXPECT().Run(gomock.Any(), gomock.Any()).Return(map[string]any{"result": "ok"}, nil) var toolset tools.Tools toolset.Init() toolset.RegisterTool(tool) a := &Agent{ ChatMessageStore: store, LLM: client, Model: "test-model", Tools: toolset, MaxIterations: 4, Session: &api.Session{ ID: "test-session", ChatMessageStore: store, AgentState: api.AgentStateIdle, }, } if err := a.Init(ctx); err != nil { t.Fatalf("init: %v", err) } if err := a.Run(ctx, ""); err != nil { t.Fatalf("run: %v", err) } // Expect prompt (UI-driven startup, no greeting message) m1 := recvMsg(t, ctx, a.Output) if m1.Type != api.MessageTypeUserInputRequest { t.Fatalf("expected user-input-request, got %v", m1.Type) } // Send a query (UI -> Agent) a.Input <- &api.UserInputResponse{Query: "test"} // Wait for choice request indicating state waiting for input. choiceMsg := recvUntil(t, ctx, a.Output, func(m *api.Message) bool { return m.Type == api.MessageTypeUserChoiceRequest }) if choiceMsg == nil { t.Fatalf("did not receive choice request") } if st := a.AgentState(); st != api.AgentStateWaitingForInput { t.Fatalf("expected waiting-for-input state, got %s", st) } // Approve tool execution (UI -> Agent) a.Input <- &api.UserChoiceResponse{Choice: 1} // Expect tool invocation messages and final response. sawToolReq, sawToolResp, sawFinal := false, false, false for !(sawToolReq && sawToolResp && sawFinal) { select { case v := <-a.Output: m, ok := v.(*api.Message) if !ok { t.Fatalf("expected *api.Message on output, got %T", v) break } switch m.Type { case api.MessageTypeToolCallRequest: sawToolReq = true case api.MessageTypeToolCallResponse: sawToolResp = true case api.MessageTypeText: if m.Source == api.MessageSourceModel { sawFinal = true } } case <-ctx.Done(): t.Fatalf("timeout before complete tool execution flow: req=%v resp=%v final=%v", sawToolReq, sawToolResp, sawFinal) } } // After final model text, the agent may either prompt for more input (UI loop) // or declare Done depending on configuration. Accept either behavior. select { case v := <-a.Output: if m, ok := v.(*api.Message); ok { if m.Type != api.MessageTypeUserInputRequest && m.Type != api.MessageTypeText { t.Fatalf("unexpected message after final model text: type=%v", m.Type) } } default: if st := a.AgentState(); st != api.AgentStateDone && st != api.AgentStateWaitingForInput { t.Fatalf("unexpected state after tool run: %s (want Done or WaitingForInput)", st) } } } func TestAgentEndToEndMetaClear(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() store := sessions.NewInMemoryChatStore() store.AddChatMessage(&api.Message{ID: "u1", Source: api.MessageSourceUser, Type: api.MessageTypeText, Payload: "hi"}) store.AddChatMessage(&api.Message{ID: "a1", Source: api.MessageSourceAgent, Type: api.MessageTypeText, Payload: "hello"}) client := mocks.NewMockClient(ctrl) chat := mocks.NewMockChat(ctrl) client.EXPECT().StartChat(gomock.Any(), "test-model").Return(chat) chat.EXPECT().Initialize(gomock.Any()).Return(nil).Times(2) // second init after clear chat.EXPECT().SetFunctionDefinitions(gomock.Any()).Return(nil) var toolset tools.Tools toolset.Init() a := &Agent{ ChatMessageStore: store, LLM: client, Model: "test-model", Tools: toolset, Session: &api.Session{ ID: "test-session", ChatMessageStore: store, AgentState: api.AgentStateIdle, }, } if err := a.Init(ctx); err != nil { t.Fatalf("init: %v", err) } if err := a.Run(ctx, ""); err != nil { t.Fatalf("run: %v", err) } // Expect startup prompt (no greeting message, UserInputRequest not stored) m1 := recvMsg(t, ctx, a.Output) if m1.Type != api.MessageTypeUserInputRequest { t.Fatalf("expected user-input-request, got %v", m1.Type) } // Only pre-seeded messages should be in store (UserInputRequest is not stored) if got := len(store.ChatMessages()); got != 2 { t.Fatalf("precondition: expected 2 messages before clear, got %d", got) } // UI sends the meta command a.Input <- &api.UserInputResponse{Query: "clear"} sawClear, sawPrompt := false, false for !(sawClear && sawPrompt) { select { case v := <-a.Output: m, ok := v.(*api.Message) if !ok { t.Fatalf("expected *api.Message on output, got %T", v) break } if m.Type == api.MessageTypeText && m.Payload == "Cleared the conversation." { sawClear = true } if sawClear && m.Type == api.MessageTypeUserInputRequest { sawPrompt = true } case <-ctx.Done(): t.Fatalf("timeout waiting for clear confirmation and prompt: %v", ctx.Err()) } } // Only the clear confirmation should be stored (UserInputRequest is not stored) msgs := store.ChatMessages() if len(msgs) != 1 { t.Fatalf("expected 1 message after clear, got %d", len(msgs)) } if msgs[0].Payload != "Cleared the conversation." { t.Fatalf("first message after clear = %q, want %q", msgs[0].Payload, "Cleared the conversation.") } } ================================================ FILE: pkg/agent/conversation.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package agent import ( "context" _ "embed" "encoding/json" "fmt" "html/template" "io" "os" "runtime" "sort" "strings" "sync" "time" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/mcp" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sessions" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" "github.com/google/uuid" "k8s.io/klog/v2" ) //go:embed systemprompt_template_default.txt var defaultSystemPromptTemplate string type Agent struct { // Input is the channel to receive user input. Input chan any // Output is the channel to send messages to the UI. Output chan any // RunOnce indicates if the agent should run only once. // If true, the agent will run only once and then exit. // If false, the agent will run in a loop until the context is done. RunOnce bool // InitialQuery is the initial query to the agent. // If provided, the agent will run only once and then exit. InitialQuery string // tool calls that are pending execution // These will typically be all the tool calls suggested by the LLM in the // previous iteration of the agentic loop. pendingFunctionCalls []ToolCallAnalysis // currChatContent tracks chat content that needs to be sent // to the LLM in the current iteration of the agentic loop. currChatContent []any // currIteration tracks the current iteration of the agentic loop. currIteration int LLM gollm.Client // PromptTemplateFile allows specifying a custom template file PromptTemplateFile string // ExtraPromptPaths allows specifying additional prompt templates // to be combined with PromptTemplateFile ExtraPromptPaths []string Model string Provider string RemoveWorkDir bool MaxIterations int // Kubeconfig is the path to the kubeconfig file. Kubeconfig string // Sandbox indicates whether to execute tools in a sandbox environment Sandbox string // SandboxImage is the container image to use for the sandbox SandboxImage string SkipPermissions bool Tools tools.Tools EnableToolUseShim bool // MCPClientEnabled indicates whether MCP client mode is enabled MCPClientEnabled bool // Recorder captures events for diagnostics Recorder journal.Recorder llmChat gollm.Chat workDir string // executor is the executor for tool execution executor sandbox.Executor // Session optionally provides a session to use. // This is used by the UI to track the state of the agent and the conversation. Session *api.Session // protects session from concurrent access sessionMu sync.Mutex // cached list of available models availableModels []string // mcpManager manages MCP client connections mcpManager *mcp.Manager // ChatMessageStore is the underlying session persistence layer. ChatMessageStore api.ChatMessageStore // SessionBackend is the configured backend for session persistence (e.g., memory, filesystem). SessionBackend string // lastErr is the most recent error run into, for use across the stack lastErr error // cancel is the function to cancel the agent's context cancel context.CancelFunc } // Assert InMemoryChatStore implements ChatMessageStore var _ api.ChatMessageStore = &sessions.InMemoryChatStore{} func (s *Agent) GetSession() *api.Session { s.sessionMu.Lock() defer s.sessionMu.Unlock() // Create a shallow copy of the session struct. The Messages slice header // is also copied, providing the caller with a snapshot of the messages // at this point in time. The UI should treat the messages as read-only // to avoid race conditions. sessionCopy := *s.Session return &sessionCopy } // addMessage creates a new message, adds it to the session, and sends it to the output channel func (c *Agent) addMessage(source api.MessageSource, messageType api.MessageType, payload any) *api.Message { c.sessionMu.Lock() defer c.sessionMu.Unlock() message := &api.Message{ ID: uuid.New().String(), Source: source, Type: messageType, Payload: payload, Timestamp: time.Now(), } // Don't store UI control signals - they're not part of the conversation if messageType != api.MessageTypeUserInputRequest { c.Session.ChatMessageStore.AddChatMessage(message) c.Session.LastModified = time.Now() } c.Output <- message return message } // setAgentState updates the agent state and ensures LastModified is updated func (c *Agent) setAgentState(newState api.AgentState) { c.sessionMu.Lock() defer c.sessionMu.Unlock() currentState := c.agentState() if currentState != newState { klog.Infof("Agent state changing from %s to %s", currentState, newState) c.Session.AgentState = newState c.Session.LastModified = time.Now() } } func (c *Agent) AgentState() api.AgentState { c.sessionMu.Lock() defer c.sessionMu.Unlock() return c.agentState() } // agentState returns the agent state without locking. // The caller is responsible for locking. func (c *Agent) agentState() api.AgentState { return c.Session.AgentState } func (s *Agent) Init(ctx context.Context) error { log := klog.FromContext(ctx) s.Input = make(chan any, 10) s.Output = make(chan any, 10) s.currIteration = 0 // when we support session, we will need to initialize this with the // current history of the conversation. s.currChatContent = []any{} if s.InitialQuery == "" && s.RunOnce { return fmt.Errorf("RunOnce mode requires an initial query to be provided") } if s.Session != nil { if s.Session.ChatMessageStore == nil { s.Session.ChatMessageStore = sessions.NewInMemoryChatStore() } s.ChatMessageStore = s.Session.ChatMessageStore if s.Session.ID == "" { s.Session.ID = uuid.New().String() } if s.Session.CreatedAt.IsZero() { s.Session.CreatedAt = time.Now() } if s.Session.LastModified.IsZero() { s.Session.LastModified = time.Now() } s.Session.Messages = s.Session.ChatMessageStore.ChatMessages() } else { return fmt.Errorf("agent requires a session to be provided") } // Create a temporary working directory workDir, err := os.MkdirTemp("", "agent-workdir-*") if err != nil { log.Error(err, "Failed to create temporary working directory") return err } log.Info("Created temporary working directory", "workDir", workDir) switch s.Sandbox { case "k8s": sandboxName := fmt.Sprintf("kubectl-ai-sandbox-%s", uuid.New().String()[:8]) // Use default image if not specified sandboxImage := s.SandboxImage if sandboxImage == "" { sandboxImage = "bitnami/kubectl:latest" } // Create sandbox with kubeconfig sb, err := sandbox.NewKubernetesSandbox(sandboxName, sandbox.WithKubeconfig(s.Kubeconfig), sandbox.WithImage(sandboxImage), ) if err != nil { return fmt.Errorf("failed to create sandbox: %w", err) } s.executor = sb log.Info("Created sandbox", "name", sandboxName, "image", sandboxImage) case "seatbelt": if runtime.GOOS != "darwin" { return fmt.Errorf("seatbelt sandbox is only supported on macOS") } s.executor = sandbox.NewSeatbeltExecutor() log.Info("Using Seatbelt executor") case "": // No sandbox, use local executor s.executor = sandbox.NewLocalExecutor() default: return fmt.Errorf("unknown sandbox type: %s", s.Sandbox) } s.workDir = workDir // Register tools with executor if none registered yet // We clone existing tools (e.g. custom tools) to ensure we have a fresh map // This avoids polluting the global default tools and ensures thread safety. s.Tools = s.Tools.CloneWithExecutor(s.executor) s.Tools.RegisterTool(tools.NewBashTool(s.executor)) s.Tools.RegisterTool(tools.NewKubectlTool(s.executor)) systemPrompt, err := s.generatePrompt(ctx, defaultSystemPromptTemplate, PromptData{ Tools: s.Tools, EnableToolUseShim: s.EnableToolUseShim, // RunOnce is a good proxy to indicate the agentic session is non-interactive mode. SessionIsInteractive: !s.RunOnce, }) if err != nil { return fmt.Errorf("generating system prompt: %w", err) } // Start a new chat session s.llmChat = gollm.NewRetryChat( s.LLM.StartChat(systemPrompt, s.Model), gollm.RetryConfig{ MaxAttempts: 3, InitialBackoff: 10 * time.Second, MaxBackoff: 60 * time.Second, BackoffFactor: 2, Jitter: true, }, ) err = s.llmChat.Initialize(s.Session.ChatMessageStore.ChatMessages()) if err != nil { return fmt.Errorf("initializing chat session: %w", err) } if s.MCPClientEnabled { if err := s.InitializeMCPClient(ctx); err != nil { klog.Errorf("Failed to initialize MCP client: %v", err) return fmt.Errorf("failed to initialize MCP client: %w", err) } // Update MCP status in session if err := s.UpdateMCPStatus(ctx, s.MCPClientEnabled); err != nil { klog.Warningf("Failed to update MCP status: %v", err) } } if !s.EnableToolUseShim { var functionDefinitions []*gollm.FunctionDefinition for _, tool := range s.Tools.AllTools() { functionDefinitions = append(functionDefinitions, tool.FunctionDefinition()) } // Sort function definitions to help KV cache reuse sort.Slice(functionDefinitions, func(i, j int) bool { return functionDefinitions[i].Name < functionDefinitions[j].Name }) if err := s.llmChat.SetFunctionDefinitions(functionDefinitions); err != nil { return fmt.Errorf("setting function definitions: %w", err) } } return nil } func (c *Agent) Close() error { if c.workDir != "" { if c.RemoveWorkDir { if err := os.RemoveAll(c.workDir); err != nil { klog.Warningf("error cleaning up directory %q: %v", c.workDir, err) } } } // Close MCP client connections if err := c.CloseMCPClient(); err != nil { klog.Warningf("error closing MCP client: %v", err) } // Close sandbox if enabled // Close executor if it exists if c.executor != nil { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() if err := c.executor.Close(ctx); err != nil { klog.Warningf("error cleaning up executor: %v", err) } else { klog.Info("Executor cleaned up successfully") } } // Cancel the agent's context if c.cancel != nil { c.cancel() } // Close the LLM client if c.LLM != nil { if err := c.LLM.Close(); err != nil { klog.Warningf("error closing LLM client: %v", err) } } return nil } func (c *Agent) LastErr() error { return c.lastErr } func (c *Agent) Run(ctx context.Context, initialQuery string) error { log := klog.FromContext(ctx) if c.Recorder != nil { ctx = journal.ContextWithRecorder(ctx, c.Recorder) } // Save unexpected error and return it in for RunOnce mode log.Info("Starting agent loop", "initialQuery", initialQuery, "runOnce", c.RunOnce) go func() { // If initialQuery is empty, try to use the one from the struct if initialQuery == "" { initialQuery = c.InitialQuery } if initialQuery != "" { c.addMessage(api.MessageSourceUser, api.MessageTypeText, initialQuery) answer, handled, err := c.handleMetaQuery(ctx, initialQuery) if err != nil { log.Error(err, "error handling meta query") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: "+err.Error()) } else if handled { // initialQuery is the 'exit' or 'quit' metaquery if c.AgentState() == api.AgentStateExited { c.addMessage(api.MessageSourceAgent, api.MessageTypeText, answer) close(c.Output) return } // we handled the meta query, so we don't need to run the agentic loop c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.addMessage(api.MessageSourceAgent, api.MessageTypeText, answer) } else { // Start the agentic loop with the initial query c.setAgentState(api.AgentStateRunning) c.currIteration = 0 c.currChatContent = []any{initialQuery} c.pendingFunctionCalls = []ToolCallAnalysis{} } } c.lastErr = nil for { var userInput any log.Info("Agent loop iteration", "state", c.AgentState()) switch c.AgentState() { case api.AgentStateIdle, api.AgentStateDone: // In RunOnce mode, we are done, so exit if c.RunOnce { log.Info("RunOnce mode, exiting agent loop") c.setAgentState(api.AgentStateExited) return } log.Info("initiating user input") c.addMessage(api.MessageSourceAgent, api.MessageTypeUserInputRequest, ">>>") select { case <-ctx.Done(): log.Info("Agent loop done") return case userInput = <-c.Input: log.Info("Received input from channel", "userInput", userInput) if userInput == io.EOF { log.Info("Agent loop done, EOF received") c.setAgentState(api.AgentStateExited) c.addMessage(api.MessageSourceAgent, api.MessageTypeText, "It has been a pleasure assisting you. Have a great day!") return } if sessionPickerResp, ok := userInput.(*api.SessionPickerResponse); ok { if sessionPickerResp.Cancelled { continue } if err := c.LoadSession(sessionPickerResp.SessionID); err != nil { log.Error(err, "error loading session") c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error loading session: "+err.Error()) } else { c.addMessage(api.MessageSourceAgent, api.MessageTypeText, fmt.Sprintf("Switched to session %s", sessionPickerResp.SessionID)) } continue } query, ok := userInput.(*api.UserInputResponse) if !ok { log.Error(nil, "Received unexpected input from channel", "userInput", userInput) return } if strings.TrimSpace(query.Query) == "" { log.Info("No query provided, skipping agentic loop") continue } c.addMessage(api.MessageSourceUser, api.MessageTypeText, query.Query) // we don't need the agentic loop for meta queries // for ex. model, tools, etc. answer, handled, err := c.handleMetaQuery(ctx, query.Query) if err != nil { log.Error(err, "error handling meta query") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: "+err.Error()) continue } if handled { // metaquery set the state to 'Exited', so we should exit if c.AgentState() == api.AgentStateExited { c.addMessage(api.MessageSourceAgent, api.MessageTypeText, answer) close(c.Output) return } // metaquery set up an interactive picker, wait for response if c.AgentState() == api.AgentStateWaitingForInput { continue } // we handled the meta query, so we don't need to run the agentic loop c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} if answer != "" { c.addMessage(api.MessageSourceAgent, api.MessageTypeText, answer) } continue } c.setAgentState(api.AgentStateRunning) c.currIteration = 0 c.currChatContent = []any{query.Query} c.pendingFunctionCalls = []ToolCallAnalysis{} log.Info("Set agent state to running, will process agentic loop", "currIteration", c.currIteration, "currChatContent", len(c.currChatContent)) } case api.AgentStateWaitingForInput: // In RunOnce mode, if we need user choice, exit with error if c.RunOnce { log.Error(nil, "RunOnce mode cannot handle user choice requests") c.setAgentState(api.AgentStateExited) c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: RunOnce mode cannot handle user choice requests") return } select { case <-ctx.Done(): log.Info("Agent loop done") return case userInput = <-c.Input: if userInput == io.EOF { log.Info("Agent loop done, EOF received") c.setAgentState(api.AgentStateExited) c.addMessage(api.MessageSourceAgent, api.MessageTypeText, "It has been a pleasure assisting you. Have a great day!") return } switch response := userInput.(type) { case *api.SessionPickerResponse: if response.Cancelled { c.setAgentState(api.AgentStateDone) continue } if err := c.LoadSession(response.SessionID); err != nil { log.Error(err, "error loading session") c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error loading session: "+err.Error()) } else { c.addMessage(api.MessageSourceAgent, api.MessageTypeText, fmt.Sprintf("Switched to session %s", response.SessionID)) } c.setAgentState(api.AgentStateDone) continue case *api.UserChoiceResponse: dispatchToolCalls := c.handleChoice(ctx, response) if dispatchToolCalls { if err := c.DispatchToolCalls(ctx); err != nil { log.Error(err, "error dispatching tool calls") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.Session.LastModified = time.Now() c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: "+err.Error()) // In RunOnce mode, exit on tool execution error if c.RunOnce { c.setAgentState(api.AgentStateExited) c.lastErr = err return } continue } // Clear pending function calls after execution c.pendingFunctionCalls = []ToolCallAnalysis{} c.setAgentState(api.AgentStateRunning) c.currIteration = c.currIteration + 1 } else { // if user has declined, we are done with this iteration c.currIteration = c.currIteration + 1 c.pendingFunctionCalls = []ToolCallAnalysis{} c.setAgentState(api.AgentStateRunning) c.Session.LastModified = time.Now() } default: log.Error(nil, "Received unexpected input from channel", "userInput", userInput) return } } case api.AgentStateRunning: // Agent is running, don't wait for input, just continue to process the agentic loop log.Info("Agent is in running state, processing agentic loop") case api.AgentStateExited: log.Info("Agent exited in RunOnce mode") return } if c.AgentState() == api.AgentStateRunning { log.Info("Processing agentic loop", "currIteration", c.currIteration, "maxIterations", c.MaxIterations, "currChatContentLen", len(c.currChatContent)) if c.currIteration >= c.MaxIterations { c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.addMessage(api.MessageSourceAgent, api.MessageTypeText, "Maximum number of iterations reached.") continue } // we run the agentic loop for one iteration stream, err := c.llmChat.SendStreaming(ctx, c.currChatContent...) if err != nil { log.Error(err, "error sending streaming LLM response") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.lastErr = err continue } // Clear our "response" now that we sent the last response c.currChatContent = nil if c.EnableToolUseShim { // convert the candidate response into a gollm.ChatResponse stream, err = candidateToShimCandidate(stream) if err != nil { c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} // In RunOnce mode, exit on shim conversion error if c.RunOnce { c.setAgentState(api.AgentStateExited) return } continue } } // Process each part of the response var functionCalls []gollm.FunctionCall // accumulator for streamed text var streamedText string var llmError error for response, err := range stream { if err != nil { log.Error(err, "error reading streaming LLM response") llmError = err c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.lastErr = llmError break } if response == nil { // end of streaming response break } // klog.Infof("response: %+v", response) if len(response.Candidates()) == 0 { llmError = fmt.Errorf("no candidates in response") log.Error(nil, "No candidates in response") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} break } candidate := response.Candidates()[0] for _, part := range candidate.Parts() { // Check if it's a text response if text, ok := part.AsText(); ok { log.Info("text response", "text", text) streamedText += text } // Check if it's a function call if calls, ok := part.AsFunctionCalls(); ok && len(calls) > 0 { log.Info("function calls", "calls", calls) functionCalls = append(functionCalls, calls...) } } } if llmError != nil { log.Error(llmError, "error streaming LLM response") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: "+llmError.Error()) c.lastErr = llmError continue } log.Info("streamedText", "streamedText", streamedText) if streamedText != "" { c.addMessage(api.MessageSourceModel, api.MessageTypeText, streamedText) } // If no function calls to be made, we're done if len(functionCalls) == 0 { log.Info("No function calls to be made, so most likely the task is completed, so we're done.") c.setAgentState(api.AgentStateDone) c.currChatContent = []any{} c.currIteration = 0 c.pendingFunctionCalls = []ToolCallAnalysis{} log.Info("Agent task completed, transitioning to done state") if streamedText == "" { // If no tool calls to be made and we do not have a response from the LLM // we should let the user know for better diagnostics. // IMPORTANT: This also prevents UIs from getting blocked on reading from the output channel. log.Info("Empty response with no tool calls from LLM.") c.addMessage(api.MessageSourceAgent, api.MessageTypeText, "Empty response from LLM") } continue } toolCallAnalysisResults, err := c.analyzeToolCalls(ctx, functionCalls) if err != nil { log.Error(err, "error analyzing tool calls") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.Session.LastModified = time.Now() c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: "+err.Error()) c.lastErr = err continue } // mark the tools for dispatching c.pendingFunctionCalls = toolCallAnalysisResults interactiveToolCallIndex := -1 modifiesResourceToolCallIndex := -1 for i, result := range toolCallAnalysisResults { if result.ModifiesResourceStr != "no" { modifiesResourceToolCallIndex = i } if result.IsInteractive { interactiveToolCallIndex = i } } if interactiveToolCallIndex >= 0 { // Show error block for both shim enabled and disabled modes errorMessage := fmt.Sprintf(" %s\n", toolCallAnalysisResults[interactiveToolCallIndex].IsInteractiveError.Error()) c.addMessage(api.MessageSourceAgent, api.MessageTypeError, errorMessage) if c.EnableToolUseShim { // Add the error as an observation observation := fmt.Sprintf("Result of running %q:\n%v", toolCallAnalysisResults[interactiveToolCallIndex].FunctionCall.Name, toolCallAnalysisResults[interactiveToolCallIndex].IsInteractiveError.Error()) c.currChatContent = append(c.currChatContent, observation) } else { // For models with tool-use support (shim disabled), use proper FunctionCallResult // Note: This assumes the model supports sending FunctionCallResult c.currChatContent = append(c.currChatContent, gollm.FunctionCallResult{ ID: toolCallAnalysisResults[interactiveToolCallIndex].FunctionCall.ID, Name: toolCallAnalysisResults[interactiveToolCallIndex].FunctionCall.Name, Result: map[string]any{"error": toolCallAnalysisResults[interactiveToolCallIndex].IsInteractiveError.Error()}, }) } c.pendingFunctionCalls = []ToolCallAnalysis{} // reset pending function calls c.currIteration = c.currIteration + 1 continue // Skip execution for interactive commands } if !c.SkipPermissions && modifiesResourceToolCallIndex >= 0 { // In RunOnce mode, exit with error if permission is required if c.RunOnce { var commandDescriptions []string for _, call := range c.pendingFunctionCalls { commandDescriptions = append(commandDescriptions, call.ParsedToolCall.Description()) } errorMessage := "RunOnce mode cannot handle permission requests. The following commands require approval:\n* " + strings.Join(commandDescriptions, "\n* ") errorMessage += "\nUse --skip-permissions flag to bypass permission checks in RunOnce mode." log.Error(nil, "RunOnce mode cannot handle permission requests", "commands", commandDescriptions) c.setAgentState(api.AgentStateExited) c.addMessage(api.MessageSourceAgent, api.MessageTypeError, errorMessage) c.lastErr = fmt.Errorf("%s", errorMessage) return } var commandDescriptions []string for _, call := range c.pendingFunctionCalls { commandDescriptions = append(commandDescriptions, call.ParsedToolCall.Description()) } confirmationPrompt := "The following commands require your approval to run:\n* " + strings.Join(commandDescriptions, "\n* ") confirmationPrompt += "\n\nDo you want to proceed ?" choiceRequest := &api.UserChoiceRequest{ Prompt: confirmationPrompt, Options: []api.UserChoiceOption{ {Value: "yes", Label: "Yes"}, {Value: "yes_and_dont_ask_me_again", Label: "Yes, and don't ask me again"}, {Value: "no", Label: "No"}, }, } c.setAgentState(api.AgentStateWaitingForInput) c.addMessage(api.MessageSourceAgent, api.MessageTypeUserChoiceRequest, choiceRequest) // Request input from the user by sending a message on the output channel. // Remaining part of the loop will be now resumed when we receive a choice input // from the user. continue } // we are here means we are in the clear to dispatch the tool calls if err := c.DispatchToolCalls(ctx); err != nil { log.Error(err, "error dispatching tool calls") c.setAgentState(api.AgentStateDone) c.pendingFunctionCalls = []ToolCallAnalysis{} c.Session.LastModified = time.Now() c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Error: "+err.Error()) c.lastErr = err continue } c.currIteration = c.currIteration + 1 c.pendingFunctionCalls = []ToolCallAnalysis{} log.Info("Tool calls dispatched successfully", "currIteration", c.currIteration, "currChatContentLen", len(c.currChatContent), "agentState", c.AgentState()) } } }() return nil } func (c *Agent) handleMetaQuery(ctx context.Context, query string) (answer string, handled bool, err error) { switch query { case "clear", "reset": c.sessionMu.Lock() // TODO: Remove this check when session persistence is default if err := c.Session.ChatMessageStore.ClearChatMessages(); err != nil { return "Failed to clear the conversation", false, err } c.llmChat.Initialize(c.Session.ChatMessageStore.ChatMessages()) c.sessionMu.Unlock() return "Cleared the conversation.", true, nil case "exit", "quit": c.setAgentState(api.AgentStateExited) return "It has been a pleasure assisting you. Have a great day!", true, nil case "model": return "Current model is `" + c.Model + "`", true, nil case "models": models, err := c.listModels(ctx) if err != nil { return "", false, fmt.Errorf("listing models: %w", err) } return "Available models:\n\n - " + strings.Join(models, "\n - ") + "\n\n", true, nil case "tools": return "Available tools:\n\n - " + strings.Join(c.Tools.Names(), "\n - ") + "\n\n", true, nil case "session": if c.SessionBackend != "filesystem" { return "Ephemeral session (memory backed). No persistent info available.", true, nil } return fmt.Sprintf("Current session:\n\n%s", c.Session.String()), true, nil case "save-session": savedSessionID, err := c.SaveSession() if err != nil { return "", false, fmt.Errorf("failed to save session: %w", err) } return "Saved session as " + savedSessionID, true, nil case "sessions": sessions, err := c.ListSessions() if err != nil { return "", false, err } if len(sessions) == 0 { return "No sessions found.", true, nil } // Add ```text so markdown doesn't wreck the format availableSessions := "```text" availableSessions += "Available sessions:\n\n" availableSessions += "ID\t\t\tCreated\t\t\tLast Accessed\t\tModel\t\tProvider\n" availableSessions += "--\t\t\t-------\t\t\t-------------\t\t-----\t\t--------\n" for _, session := range sessions { availableSessions += fmt.Sprintf("%s\t%s\t%s\t%s\t%s\n", session.ID, session.CreatedAt.Format("2006-01-02 15:04"), session.LastModified.Format("2006-01-02 15:04"), session.ModelID, session.ProviderID) } // close the ```text box availableSessions += "```" return availableSessions, true, nil } if strings.HasPrefix(query, "resume-session") { parts := strings.Split(query, " ") if len(parts) != 2 { return "Invalid command. Usage: resume-session ", true, nil } sessionID := parts[1] if err := c.LoadSession(sessionID); err != nil { return "", false, err } return fmt.Sprintf("Resumed session %s.", sessionID), true, nil } return "", false, nil } func (c *Agent) NewSession() (string, error) { if _, err := c.SaveSession(); err != nil { return "", fmt.Errorf("failed to save current session: %w", err) } manager, err := sessions.NewSessionManager(c.SessionBackend) if err != nil { return "", fmt.Errorf("failed to create session manager: %w", err) } metadata := sessions.Metadata{ ModelID: c.Model, ProviderID: c.Provider, } newSession, err := manager.NewSession(metadata) if err != nil { return "", fmt.Errorf("failed to create new session: %w", err) } // If we are using a sandbox, we should spin up a new one for the new session if c.Sandbox == "k8s" { sandboxName := fmt.Sprintf("kubectl-ai-sandbox-%s", uuid.New().String()[:8]) sandboxImage := c.SandboxImage sb, err := sandbox.NewKubernetesSandbox(sandboxName, sandbox.WithKubeconfig(c.Kubeconfig), sandbox.WithImage(sandboxImage), ) if err != nil { return "", fmt.Errorf("failed to create new sandbox: %w", err) } c.sessionMu.Lock() if c.executor != nil { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) if err := c.executor.Close(ctx); err != nil { klog.Warningf("error closing old executor: %v", err) } cancel() } c.executor = sb klog.Info("Created new sandbox for new session", "name", sandboxName) // Re-bind all tools to the new executor c.Tools = c.Tools.CloneWithExecutor(c.executor) c.Tools.RegisterTool(tools.NewBashTool(c.executor)) c.Tools.RegisterTool(tools.NewKubectlTool(c.executor)) c.sessionMu.Unlock() } if err := c.LoadSession(newSession.ID); err != nil { return "", fmt.Errorf("failed to load new session: %w", err) } return newSession.ID, nil } func (c *Agent) SaveSession() (string, error) { c.sessionMu.Lock() defer c.sessionMu.Unlock() manager, err := sessions.NewSessionManager(c.SessionBackend) if err != nil { return "", fmt.Errorf("failed to create session manager: %w", err) } if c.Session != nil { foundSession, _ := manager.FindSessionByID(c.Session.ID) if foundSession != nil { return foundSession.ID, nil } } metadata := sessions.Metadata{ CreatedAt: c.Session.CreatedAt, LastAccessed: time.Now(), ModelID: c.Model, ProviderID: c.Provider, } newSession, err := manager.NewSession(metadata) if err != nil { return "", fmt.Errorf("failed to create new session: %w", err) } messages := c.ChatMessageStore.ChatMessages() if err := newSession.ChatMessageStore.SetChatMessages(messages); err != nil { return "", fmt.Errorf("failed to save chat messages to new session: %w", err) } c.ChatMessageStore = newSession.ChatMessageStore c.Session = newSession c.Session.Messages = messages if c.llmChat != nil { _ = c.llmChat.Initialize(c.Session.ChatMessageStore.ChatMessages()) } return newSession.ID, nil } // LoadSession loads a session by ID (or latest), updates the agent's state, and re-initializes the chat. func (c *Agent) LoadSession(sessionID string) error { manager, err := sessions.NewSessionManager(c.SessionBackend) if err != nil { return fmt.Errorf("failed to create session manager: %w", err) } var session *api.Session if sessionID == "" || sessionID == "latest" { s, err := manager.GetLatestSession() if err != nil { return fmt.Errorf("failed to get latest session: %w", err) } if s == nil { return fmt.Errorf("no sessions found to resume") } session = s } else { s, err := manager.FindSessionByID(sessionID) if err != nil { return fmt.Errorf("failed to get session %q: %w", sessionID, err) } session = s } c.sessionMu.Lock() defer c.sessionMu.Unlock() if session.ChatMessageStore == nil { session.ChatMessageStore = sessions.NewInMemoryChatStore() } c.Session = session c.ChatMessageStore = session.ChatMessageStore c.Session.Messages = session.ChatMessageStore.ChatMessages() c.Session.LastModified = time.Now() // Reset state if it was left running (e.g. from a crash) if c.Session.AgentState == api.AgentStateRunning || c.Session.AgentState == api.AgentStateInitializing { c.Session.AgentState = api.AgentStateIdle } if err := manager.UpdateLastAccessed(session); err != nil { return fmt.Errorf("failed to update session metadata: %w", err) } if c.llmChat != nil { if err := c.llmChat.Initialize(c.Session.ChatMessageStore.ChatMessages()); err != nil { return fmt.Errorf("failed to re-initialize chat with new session: %w", err) } } return nil } // ListSessions returns available sessions for UI pickers func (c *Agent) ListSessions() ([]api.SessionInfo, error) { manager, err := sessions.NewSessionManager(c.SessionBackend) if err != nil { return nil, fmt.Errorf("failed to create session manager: %w", err) } sessionList, err := manager.ListSessions() if err != nil { return nil, fmt.Errorf("failed to list sessions: %w", err) } sessionInfos := make([]api.SessionInfo, len(sessionList)) for i, session := range sessionList { msgCount := 0 if session.ChatMessageStore != nil { msgCount = len(session.ChatMessageStore.ChatMessages()) } sessionInfos[i] = api.SessionInfo{ ID: session.ID, Name: session.Name, ModelID: session.ModelID, ProviderID: session.ProviderID, CreatedAt: session.CreatedAt, LastModified: session.LastModified, MessageCount: msgCount, } } return sessionInfos, nil } func (c *Agent) listModels(ctx context.Context) ([]string, error) { if c.availableModels == nil { modelNames, err := c.LLM.ListModels(ctx) if err != nil { return nil, fmt.Errorf("listing models: %w", err) } c.availableModels = modelNames } return c.availableModels, nil } func (c *Agent) DispatchToolCalls(ctx context.Context) error { log := klog.FromContext(ctx) // execute all pending function calls for _, call := range c.pendingFunctionCalls { // Only show "Running" message and proceed with execution for non-interactive commands toolDescription := call.ParsedToolCall.Description() c.addMessage(api.MessageSourceModel, api.MessageTypeToolCallRequest, toolDescription) output, err := call.ParsedToolCall.InvokeTool(ctx, tools.InvokeToolOptions{ Kubeconfig: c.Kubeconfig, WorkDir: c.workDir, Executor: c.executor, }) if err != nil { log.Error(err, "error executing action", "output", output) c.addMessage(api.MessageSourceAgent, api.MessageTypeToolCallResponse, err.Error()) return err } // Handle timeout message using UI blocks if execResult, ok := output.(*sandbox.ExecResult); ok && execResult != nil && execResult.StreamType == "timeout" { c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "\nTimeout reached after 7 seconds\n") } // Add the tool call result to maintain conversation flow var payload any if c.EnableToolUseShim { // Add the error as an observation observation := fmt.Sprintf("Result of running %q:\n%v", call.FunctionCall.Name, output) c.currChatContent = append(c.currChatContent, observation) payload = observation } else { // If shim is disabled, convert the result to a map and append FunctionCallResult result, err := tools.ToolResultToMap(output) if err != nil { log.Error(err, "error converting tool result to map", "output", output) return err } payload = result c.currChatContent = append(c.currChatContent, gollm.FunctionCallResult{ ID: call.FunctionCall.ID, Name: call.FunctionCall.Name, Result: result, }) } c.addMessage(api.MessageSourceAgent, api.MessageTypeToolCallResponse, payload) } return nil } // The key idea is to treat all tool calls to be executed atomically or not // If all tool calls are readonly call, it is straight forward // if some of the tool calls are not readonly, then the interesting question is should the permission // be asked for each of the tool call or only once for all the tool calls. // I think treating all tool calls as atomic is the right thing to do. type ToolCallAnalysis struct { FunctionCall gollm.FunctionCall ParsedToolCall *tools.ToolCall IsInteractive bool IsInteractiveError error ModifiesResourceStr string } func (c *Agent) analyzeToolCalls(ctx context.Context, toolCalls []gollm.FunctionCall) ([]ToolCallAnalysis, error) { toolCallAnalysis := make([]ToolCallAnalysis, len(toolCalls)) for i, call := range toolCalls { toolCallAnalysis[i].FunctionCall = call toolCall, err := c.Tools.ParseToolInvocation(ctx, call.Name, call.Arguments) if err != nil { return nil, fmt.Errorf("error parsing tool call: %w", err) } toolCallAnalysis[i].IsInteractive, err = toolCall.GetTool().IsInteractive(call.Arguments) if err != nil { toolCallAnalysis[i].IsInteractiveError = err } toolCallAnalysis[i].ModifiesResourceStr = toolCall.GetTool().CheckModifiesResource(call.Arguments) toolCallAnalysis[i].ParsedToolCall = toolCall } return toolCallAnalysis, nil } func (c *Agent) handleChoice(ctx context.Context, choice *api.UserChoiceResponse) (dispatchToolCalls bool) { log := klog.FromContext(ctx) // if user input is a choice and use has declined the operation, // we need to abort all pending function calls. // update the currChatContent with the choice and keep the agent loop running. // Normalize the input switch choice.Choice { case 1: dispatchToolCalls = true case 2: c.SkipPermissions = true dispatchToolCalls = true case 3: c.currChatContent = append(c.currChatContent, gollm.FunctionCallResult{ ID: c.pendingFunctionCalls[0].FunctionCall.ID, Name: c.pendingFunctionCalls[0].FunctionCall.Name, Result: map[string]any{ "error": "User declined to run this operation.", "status": "declined", "retryable": false, }, }) c.pendingFunctionCalls = []ToolCallAnalysis{} dispatchToolCalls = false c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Operation was skipped. User declined to run this operation.") default: // This case should technically not be reachable due to AskForConfirmation loop err := fmt.Errorf("invalid confirmation choice: %q", choice.Choice) log.Error(err, "Invalid choice received from AskForConfirmation") c.pendingFunctionCalls = []ToolCallAnalysis{} dispatchToolCalls = false c.addMessage(api.MessageSourceAgent, api.MessageTypeError, "Invalid choice received. Cancelling operation.") } return dispatchToolCalls } // generateFromTemplate generates a prompt for LLM. It uses the prompt from the provides template file or default. func (a *Agent) generatePrompt(_ context.Context, defaultPromptTemplate string, data PromptData) (string, error) { promptTemplate := defaultPromptTemplate if a.PromptTemplateFile != "" { content, err := os.ReadFile(a.PromptTemplateFile) if err != nil { return "", fmt.Errorf("error reading template file: %v", err) } promptTemplate = string(content) } for _, extraPromptPath := range a.ExtraPromptPaths { content, err := os.ReadFile(extraPromptPath) if err != nil { return "", fmt.Errorf("error reading extra prompt path: %v", err) } promptTemplate += "\n" + string(content) } tmpl, err := template.New("promptTemplate").Parse(promptTemplate) if err != nil { return "", fmt.Errorf("building template for prompt: %w", err) } var result strings.Builder err = tmpl.Execute(&result, &data) if err != nil { return "", fmt.Errorf("evaluating template for prompt: %w", err) } return result.String(), nil } // PromptData represents the structure of the data to be filled into the template. type PromptData struct { Query string Tools tools.Tools EnableToolUseShim bool SessionIsInteractive bool } func (a *PromptData) ToolsAsJSON() string { var toolDefinitions []*gollm.FunctionDefinition for _, tool := range a.Tools.AllTools() { toolDefinitions = append(toolDefinitions, tool.FunctionDefinition()) } json, err := json.MarshalIndent(toolDefinitions, "", " ") if err != nil { return "" } return string(json) } func (a *PromptData) ToolNames() string { return strings.Join(a.Tools.Names(), ", ") } type ReActResponse struct { Thought string `json:"thought"` Answer string `json:"answer,omitempty"` Action *Action `json:"action,omitempty"` } type Action struct { Name string `json:"name"` Reason string `json:"reason"` Command string `json:"command"` ModifiesResource string `json:"modifies_resource"` } func extractJSON(s string) (string, bool) { const jsonBlockMarker = "```json" first := strings.Index(s, jsonBlockMarker) last := strings.LastIndex(s, "```") if first == -1 || last == -1 || first == last { return "", false } data := s[first+len(jsonBlockMarker) : last] return data, true } // parseReActResponse parses the LLM response into a ReActResponse struct // This function assumes the input contains exactly one JSON code block // formatted with ```json and ``` markers. The JSON block is expected to // contain a valid ReActResponse object. func parseReActResponse(input string) (*ReActResponse, error) { cleaned, found := extractJSON(input) if !found { return nil, fmt.Errorf("no JSON code block found in %q", cleaned) } cleaned = strings.ReplaceAll(cleaned, "\n", "") cleaned = strings.TrimSpace(cleaned) var reActResp ReActResponse if err := json.Unmarshal([]byte(cleaned), &reActResp); err != nil { return nil, fmt.Errorf("parsing JSON %q: %w", cleaned, err) } return &reActResp, nil } // toMap converts the value to a map, going via JSON func toMap(v any) (map[string]any, error) { j, err := json.Marshal(v) if err != nil { return nil, fmt.Errorf("converting %T to json: %w", v, err) } m := make(map[string]any) if err := json.Unmarshal(j, &m); err != nil { return nil, fmt.Errorf("converting json to map: %w", err) } return m, nil } func candidateToShimCandidate(iterator gollm.ChatResponseIterator) (gollm.ChatResponseIterator, error) { return func(yield func(gollm.ChatResponse, error) bool) { buffer := "" for response, err := range iterator { if err != nil { yield(nil, err) return } if len(response.Candidates()) == 0 { yield(nil, fmt.Errorf("no candidates in LLM response")) return } candidate := response.Candidates()[0] for _, part := range candidate.Parts() { if text, ok := part.AsText(); ok { buffer += text klog.Infof("text is %q", text) } else { yield(nil, fmt.Errorf("no text part found in candidate")) return } } } if buffer == "" { yield(nil, nil) return } parsedReActResp, err := parseReActResponse(buffer) if err != nil { yield(nil, fmt.Errorf("parsing ReAct response %q: %w", buffer, err)) return } buffer = "" // TODO: any trailing text? yield(&ShimResponse{candidate: parsedReActResp}, nil) }, nil } type ShimResponse struct { candidate *ReActResponse } func (r *ShimResponse) UsageMetadata() any { return nil } func (r *ShimResponse) Candidates() []gollm.Candidate { return []gollm.Candidate{&ShimCandidate{candidate: r.candidate}} } type ShimCandidate struct { candidate *ReActResponse } func (c *ShimCandidate) String() string { return fmt.Sprintf("Thought: %s\nAnswer: %s\nAction: %s", c.candidate.Thought, c.candidate.Answer, c.candidate.Action) } func (c *ShimCandidate) Parts() []gollm.Part { var parts []gollm.Part if c.candidate.Thought != "" { parts = append(parts, &ShimPart{text: c.candidate.Thought}) } if c.candidate.Answer != "" { parts = append(parts, &ShimPart{text: c.candidate.Answer}) } if c.candidate.Action != nil { parts = append(parts, &ShimPart{action: c.candidate.Action}) } return parts } type ShimPart struct { text string action *Action } func (p *ShimPart) AsText() (string, bool) { return p.text, p.text != "" } func (p *ShimPart) AsFunctionCalls() ([]gollm.FunctionCall, bool) { if p.action != nil { functionCallArgs, err := toMap(p.action) if err != nil { return nil, false } delete(functionCallArgs, "name") // passed separately // delete(functionCallArgs, "reason") // delete(functionCallArgs, "modifies_resource") return []gollm.FunctionCall{ { Name: p.action.Name, Arguments: functionCallArgs, }, }, true } return nil, false } ================================================ FILE: pkg/agent/conversation_test.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package agent import ( "context" "os" "strings" "testing" "time" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/internal/mocks" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sessions" "go.uber.org/mock/gomock" ) func TestHandleMetaQuery(t *testing.T) { ctx := context.Background() tests := []struct { name string query string expectations func(t *testing.T) *Agent verify func(t *testing.T, a *Agent, answer string) expect string }{ { name: "clear (shows store before/after with mocked model + tool outputs)", query: "clear", expect: "Cleared the conversation.", expectations: func(t *testing.T) *Agent { ctrl := gomock.NewController(t) t.Cleanup(ctrl.Finish) store := sessions.NewInMemoryChatStore() chat := mocks.NewMockChat(ctrl) chat.EXPECT().Initialize([]*api.Message{}).Times(1) mt := mocks.NewMockTool(ctrl) mt.EXPECT().Name().Return("mock namespace tool").AnyTimes() mt.EXPECT().FunctionDefinition().Return(&gollm.FunctionDefinition{ Name: "mock namespace tool", Description: "Inspect current Kubernetes namespace", }).AnyTimes() const toolResult = `{"namespace":"test-namespace"}` mt.EXPECT().Run(gomock.Any(), gomock.Any()). Return(toolResult, nil).Times(1) const modelText = "The current namespace is test-namespace." // user message _ = store.AddChatMessage(&api.Message{ ID: "u1", Source: api.MessageSourceUser, Type: api.MessageTypeText, Payload: "What's my current namespace?", }) // model response _ = store.AddChatMessage(&api.Message{ ID: "a1", Source: api.MessageSourceAgent, Type: api.MessageTypeText, Payload: modelText, }) // tool call result if out, err := mt.Run(ctx, map[string]any{}); err == nil { _ = store.AddChatMessage(&api.Message{ ID: "t1", Source: api.MessageSourceAgent, Type: api.MessageTypeText, Payload: out, }) } else { t.Fatalf("mock tool run failed: %v", err) } if got := len(store.ChatMessages()); got != 3 { t.Fatalf("precondition: expected 3 messages before clear, got %d", got) } a := &Agent{llmChat: chat} a.Session = &api.Session{ChatMessageStore: store} return a }, verify: func(t *testing.T, a *Agent, _ string) { if got := len(a.Session.ChatMessageStore.ChatMessages()); got != 0 { t.Fatalf("expected store to be empty after clear, got %d", got) } }, }, { name: "exit", query: "exit", expect: "It has been a pleasure assisting you. Have a great day!", expectations: func(t *testing.T) *Agent { a := &Agent{} a.Session = &api.Session{} return a }, verify: func(t *testing.T, a *Agent, _ string) { if a.AgentState() != api.AgentStateExited { t.Fatalf("expected agent to exit") } }, }, { name: "model", query: "model", expect: "Current model is `test-model`", expectations: func(t *testing.T) *Agent { a := &Agent{Model: "test-model"} a.Session = &api.Session{} return a }, }, { name: "models", query: "models", expect: "Available models:\n\n - a\n - b\n\n", expectations: func(t *testing.T) *Agent { ctrl := gomock.NewController(t) t.Cleanup(ctrl.Finish) llm := mocks.NewMockClient(ctrl) llm.EXPECT().ListModels(ctx).Return([]string{"a", "b"}, nil) a := &Agent{LLM: llm} a.Session = &api.Session{} return a }, }, { name: "tools", query: "tools", expect: "Available tools:", expectations: func(t *testing.T) *Agent { ctrl := gomock.NewController(t) t.Cleanup(ctrl.Finish) mt := mocks.NewMockTool(ctrl) mt.EXPECT().Name().Return("mocktool").AnyTimes() mt.EXPECT().FunctionDefinition().Return(&gollm.FunctionDefinition{ Name: "mocktool", Description: "Mocked tool for tests", }).AnyTimes() a := &Agent{} a.Tools.Init() a.Tools.RegisterTool(mt) a.Session = &api.Session{} return a }, verify: func(t *testing.T, _ *Agent, answer string) { if !strings.Contains(answer, "mocktool") { t.Fatalf("expected kubectl tool in output: %q", answer) } }, }, { name: "session", query: "session", expect: "Session ID:", expectations: func(t *testing.T) *Agent { oldHome := os.Getenv("HOME") t.Cleanup(func() { os.Setenv("HOME", oldHome) }) home := t.TempDir() os.Setenv("HOME", home) manager, err := sessions.NewSessionManager("memory") if err != nil { t.Fatalf("creating session manager: %v", err) } sess, err := manager.NewSession(sessions.Metadata{ProviderID: "p", ModelID: "m"}) if err != nil { t.Fatalf("creating session: %v", err) } a := &Agent{ChatMessageStore: sess.ChatMessageStore, SessionBackend: "filesystem"} a.Session = sess return a }, verify: func(t *testing.T, _ *Agent, answer string) { if !strings.Contains(answer, "ID:") { t.Fatalf("expected session info, got %q", answer) } }, }, { name: "sessions", query: "sessions", expect: "Available sessions:", expectations: func(t *testing.T) *Agent { oldHome := os.Getenv("HOME") t.Cleanup(func() { os.Setenv("HOME", oldHome) }) home := t.TempDir() os.Setenv("HOME", home) manager, err := sessions.NewSessionManager("memory") if err != nil { t.Fatalf("creating session manager: %v", err) } if _, err := manager.NewSession(sessions.Metadata{ProviderID: "p1", ModelID: "m1"}); err != nil { t.Fatalf("creating session: %v", err) } a := &Agent{SessionBackend: "memory"} a.Session = &api.Session{ChatMessageStore: sessions.NewInMemoryChatStore()} return a }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { a := tt.expectations(t) ans, handled, err := a.handleMetaQuery(ctx, tt.query) if err != nil { t.Fatalf("handleMetaQuery returned error: %v", err) } if !handled { t.Fatalf("expected query %q to be handled", tt.query) } if tt.expect != "" && !strings.Contains(ans, tt.expect) { t.Fatalf("expected %q to contain %q", ans, tt.expect) } if tt.verify != nil { tt.verify(t, a, ans) } }) } } func TestAgent_NewSession(t *testing.T) { // Setup manager, err := sessions.NewSessionManager("memory") if err != nil { t.Fatalf("creating session manager: %v", err) } // Create initial session sess1, err := manager.NewSession(sessions.Metadata{}) if err != nil { t.Fatalf("creating session 1: %v", err) } a := &Agent{ SessionBackend: "memory", } a.Session = sess1 // Call NewSession newID, err := a.NewSession() if err != nil { t.Fatalf("NewSession failed: %v", err) } if newID == sess1.ID { t.Fatalf("expected new session ID to be different from old one") } if a.Session.ID != newID { t.Fatalf("agent session ID mismatch: got %s, want %s", a.Session.ID, newID) } } func TestAgent_LoadSession_ResetsState(t *testing.T) { // Setup manager, err := sessions.NewSessionManager("memory") if err != nil { t.Fatalf("creating session manager: %v", err) } // Create a session in "running" state sess1, err := manager.NewSession(sessions.Metadata{}) if err != nil { t.Fatalf("creating session 1: %v", err) } sess1.AgentState = api.AgentStateRunning if err := manager.UpdateLastAccessed(sess1); err != nil { t.Fatalf("updating session: %v", err) } a := &Agent{ SessionBackend: "memory", } // Load the session if err := a.LoadSession(sess1.ID); err != nil { t.Fatalf("LoadSession failed: %v", err) } // Verify state is reset to idle if a.Session.AgentState != api.AgentStateIdle { t.Errorf("expected agent state to be idle, got %s", a.Session.AgentState) } } func TestAgent_Init_CreatesSessionInStore(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mocks.NewMockClient(ctrl) mockChat := mocks.NewMockChat(ctrl) // Expect StartChat to be called mockClient.EXPECT().StartChat(gomock.Any(), gomock.Any()).Return(mockChat) // Expect Initialize to be called mockChat.EXPECT().Initialize(gomock.Any()).Return(nil) // Expect SetFunctionDefinitions to be called mockChat.EXPECT().SetFunctionDefinitions(gomock.Any()).Return(nil) // Setup session := &api.Session{ ID: "test-session", AgentState: api.AgentStateIdle, ChatMessageStore: sessions.NewInMemoryChatStore(), } a := &Agent{ SessionBackend: "memory", // Init requires these Input: make(chan any), Output: make(chan any), LLM: mockClient, Session: session, } if err := a.Init(context.Background()); err != nil { t.Fatalf("Init failed: %v", err) } if a.Session != session { t.Errorf("expected agent to use provided session") } } func TestAgent_NewSession_NoDeadlock(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mocks.NewMockClient(ctrl) mockChat := mocks.NewMockChat(ctrl) // Expect StartChat to be called for initial session only mockClient.EXPECT().StartChat(gomock.Any(), gomock.Any()).Return(mockChat).Times(1) // Expect Initialize to be called for initial session AND new session (and maybe more?) mockChat.EXPECT().Initialize(gomock.Any()).Return(nil).AnyTimes() // Expect SetFunctionDefinitions to be called for initial session only mockChat.EXPECT().SetFunctionDefinitions(gomock.Any()).Return(nil).Times(1) // Setup session := &api.Session{ ID: "initial-session", AgentState: api.AgentStateIdle, ChatMessageStore: sessions.NewInMemoryChatStore(), } a := &Agent{ SessionBackend: "memory", Input: make(chan any), Output: make(chan any), LLM: mockClient, Session: session, } // Init if err := a.Init(context.Background()); err != nil { t.Fatalf("Init failed: %v", err) } // Create new session // This should not deadlock done := make(chan struct{}) go func() { if _, err := a.NewSession(); err != nil { t.Errorf("NewSession failed: %v", err) } close(done) }() select { case <-done: // Success case <-time.After(2 * time.Second): t.Fatal("NewSession timed out (potential deadlock)") } } ================================================ FILE: pkg/agent/manager.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package agent import ( "context" "fmt" "sync" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sessions" "k8s.io/klog/v2" ) // Factory is a function that creates a new Agent instance. type Factory func(context.Context) (*Agent, error) // AgentManager manages the lifecycle of agents and their sessions. type AgentManager struct { factory Factory sessionManager *sessions.SessionManager agents map[string]*Agent // sessionID -> agent mu sync.RWMutex onAgentCreated func(*Agent) } // NewAgentManager creates a new Manager. func NewAgentManager(factory Factory, sessionManager *sessions.SessionManager) *AgentManager { return &AgentManager{ factory: factory, sessionManager: sessionManager, agents: make(map[string]*Agent), } } // SetAgentCreatedCallback sets the callback to be called when a new agent is created. // It also calls the callback immediately for all currently active agents. func (sm *AgentManager) SetAgentCreatedCallback(cb func(*Agent)) { sm.mu.Lock() defer sm.mu.Unlock() sm.onAgentCreated = cb for _, agent := range sm.agents { cb(agent) } } // GetAgent returns the agent for the given session ID, loading it if necessary. func (sm *AgentManager) GetAgent(ctx context.Context, sessionID string) (*Agent, error) { sm.mu.RLock() agent, ok := sm.agents[sessionID] sm.mu.RUnlock() if ok { return agent, nil } session, err := sm.sessionManager.FindSessionByID(sessionID) if err != nil { return nil, fmt.Errorf("session not found: %w", err) } newAgent, err := sm.factory(ctx) if err != nil { return nil, fmt.Errorf("creating agent: %w", err) } return sm.startAgent(ctx, session, newAgent) } // Close closes all active agents. func (sm *AgentManager) Close() error { sm.mu.Lock() defer sm.mu.Unlock() for id, agent := range sm.agents { klog.Infof("Closing agent for session %s", id) if err := agent.Close(); err != nil { klog.Errorf("Error closing agent %s: %v", id, err) } } // Clear the map sm.agents = make(map[string]*Agent) return nil } // ListSessions delegates to the underlying store. func (sm *AgentManager) ListSessions() ([]*api.Session, error) { return sm.sessionManager.ListSessions() } // FindSessionByID delegates to the underlying store. func (sm *AgentManager) FindSessionByID(id string) (*api.Session, error) { return sm.sessionManager.FindSessionByID(id) } // DeleteSession delegates to the underlying store and closes the active agent if any. func (sm *AgentManager) DeleteSession(id string) error { sm.mu.Lock() if agent, ok := sm.agents[id]; ok { agent.Close() delete(sm.agents, id) } sm.mu.Unlock() return sm.sessionManager.DeleteSession(id) } // UpdateLastAccessed delegates to the underlying store. func (sm *AgentManager) UpdateLastAccessed(session *api.Session) error { return sm.sessionManager.UpdateLastAccessed(session) } func (sm *AgentManager) startAgent(ctx context.Context, session *api.Session, agent *Agent) (*Agent, error) { agent.Session = session if err := agent.Init(ctx); err != nil { return nil, fmt.Errorf("initializing agent: %w", err) } agentCtx, cancel := context.WithCancel(context.Background()) agent.cancel = cancel if err := agent.Run(agentCtx, ""); err != nil { cancel() return nil, fmt.Errorf("starting agent loop: %w", err) } sm.mu.Lock() sm.agents[session.ID] = agent if sm.onAgentCreated != nil { sm.onAgentCreated(agent) } sm.mu.Unlock() return agent, nil } ================================================ FILE: pkg/agent/mcp_client.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package agent import ( "context" "fmt" "strings" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/mcp" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" "k8s.io/klog/v2" ) // InitializeMCPClient initializes MCP client functionality for the agent. // It connects to servers and registers discovered tools with the kubectl-ai tool system. func (a *Agent) InitializeMCPClient(ctx context.Context) error { // Initialize the MCP manager manager, err := mcp.InitializeManager() if err != nil { return fmt.Errorf("failed to initialize MCP manager: %w", err) } // Connect to servers and register tools err = manager.RegisterWithToolSystem(ctx, func(serverName string, toolInfo mcp.Tool) error { // Create schema for the tool schema, err := tools.ConvertToolToGollm(&toolInfo) if err != nil { return err } // Create an MCPTool wrapper first to get the unique name mcpTool := tools.NewMCPTool(serverName, toolInfo.Name, toolInfo.Description, schema, manager) // Update schema with unique name and better description to avoid conflicts schema.Name = mcpTool.UniqueToolName() schema.Description = fmt.Sprintf("%s (from %s)", toolInfo.Description, serverName) // Create and register MCP tool wrapper tools.RegisterTool(mcpTool) return nil }) if err != nil { return fmt.Errorf("failed to register MCP tools: %w", err) } // Store the manager for later use a.mcpManager = manager return nil } // UpdateMCPStatus updates the MCP status in the agent's session func (a *Agent) UpdateMCPStatus(ctx context.Context, mcpClientEnabled bool) error { if a.mcpManager == nil && !mcpClientEnabled { // No MCP functionality requested return nil } status, err := a.getMCPStatus(ctx, mcpClientEnabled) if err != nil { klog.Errorf("Failed to get MCP server status: %v", err) return err } // Update the session with MCP status a.sessionMu.Lock() defer a.sessionMu.Unlock() a.Session.MCPStatus = status return nil } // getMCPStatus retrieves the current MCP status func (a *Agent) getMCPStatus(ctx context.Context, mcpClientEnabled bool) (*api.MCPStatus, error) { var mcpStatus *mcp.MCPStatus var err error if mcpClientEnabled && a.mcpManager != nil { // In client mode, use the provided manager mcpStatus, err = a.mcpManager.GetStatus(ctx, mcpClientEnabled) if err != nil { return nil, err } } else { // Create minimal status mcpStatus = &mcp.MCPStatus{ ClientEnabled: mcpClientEnabled, } } // Convert from mcp.MCPStatus to api.MCPStatus return a.convertMCPStatus(mcpStatus), nil } // convertMCPStatus converts from mcp.MCPStatus to api.MCPStatus func (a *Agent) convertMCPStatus(mcpStatus *mcp.MCPStatus) *api.MCPStatus { if mcpStatus == nil { return nil } apiStatus := &api.MCPStatus{ TotalServers: mcpStatus.TotalServers, ConnectedCount: mcpStatus.ConnectedCount, FailedCount: mcpStatus.FailedCount, TotalTools: mcpStatus.TotalTools, ClientEnabled: mcpStatus.ClientEnabled, } // Convert server connection info for _, server := range mcpStatus.ServerInfoList { apiServerInfo := api.ServerConnectionInfo{ Name: server.Name, Command: server.Command, IsLegacy: server.IsLegacy, IsConnected: server.IsConnected, } // Convert tools for _, tool := range server.AvailableTools { apiTool := api.MCPTool{ Name: tool.Name, Description: tool.Description, Server: tool.Server, } apiServerInfo.AvailableTools = append(apiServerInfo.AvailableTools, apiTool) } apiStatus.ServerInfoList = append(apiStatus.ServerInfoList, apiServerInfo) } return apiStatus } // GetMCPStatusText returns a formatted text representation of the MCP status // This can be used by UIs that want to display the status as text func (a *Agent) GetMCPStatusText() string { a.sessionMu.Lock() defer a.sessionMu.Unlock() if a.Session.MCPStatus == nil { return "" } var statusText strings.Builder status := a.Session.MCPStatus // Add summary text if status.ClientEnabled && status.ConnectedCount > 0 { statusText.WriteString(fmt.Sprintf("Successfully connected to %d MCP server(s) (%d tools discovered)\n\n", status.ConnectedCount, status.TotalTools)) } else if status.ClientEnabled { statusText.WriteString("No MCP servers connected\n\n") } else if status.TotalServers > 0 { statusText.WriteString(fmt.Sprintf("%d MCP servers configured (client mode disabled)\n\n", status.TotalServers)) } else { statusText.WriteString("No MCP servers configured\n\n") } // Add server details for _, server := range status.ServerInfoList { connectionStatus := "Disconnected" if server.IsConnected { connectionStatus = "Connected" } // Get tool names if available var toolNames []string for _, tool := range server.AvailableTools { toolNames = append(toolNames, tool.Name) } // Format server details statusText.WriteString(" • ") // Bullet point with indentation statusText.WriteString(fmt.Sprintf("%s (%s) - %s", server.Name, extractCommandName(server.Command), connectionStatus)) if len(toolNames) > 0 { statusText.WriteString(fmt.Sprintf(", Tools: %s", strings.Join(toolNames, ", "))) } statusText.WriteString("\n") } return statusText.String() } // extractCommandName gets the base command from a command string func extractCommandName(command string) string { if command == "" { return "remote" // Return 'remote' for HTTP-based servers } parts := strings.Fields(command) if len(parts) > 0 { return parts[0] } return command } // CloseMCPClient closes the MCP client connections func (a *Agent) CloseMCPClient() error { if a.mcpManager != nil { err := a.mcpManager.Close() a.mcpManager = nil return err } return nil } ================================================ FILE: pkg/agent/systemprompt_template_default.txt ================================================ You are `kubectl-ai`, an AI assistant with expertise in operating and performing actions against a kubernetes cluster. Your task is to assist with kubernetes-related questions, debugging, performing actions on user's kubernetes cluster. {{if .EnableToolUseShim }} ## Available tools {{.ToolsAsJSON}} ## Instructions: 1. Analyze the query, previous reasoning steps, and observations. 2. Reflect on 5-7 different ways to solve the given query or task. Think carefully about each solution before picking the best one. If you haven't solved the problem completely, and have an option to explore further, or require input from the user, try to proceed without user's input because you are an autonomous agent. 3. Decide on the next action: use a tool or provide a final answer and respond in the following JSON format: If you need to use a tool: ```json { "thought": "Your detailed reasoning about what to do next", "action": { "name": "Tool name ({{.ToolNames}})", "reason": "Explanation of why you chose this tool (not more than 100 words)", "command": "Complete command to be executed. For example, 'kubectl get pods', 'kubectl get ns'", "modifies_resource": "Whether the command modifies a kubernetes resource. Possible values are 'yes' or 'no' or 'unknown'" } } ``` If you have enough information to answer the query: ```json { "thought": "Your final reasoning process", "answer": "Your comprehensive answer to the query" } ``` {{else}} ## Instructions: - Examine current state of kubernetes resources relevant to user's query. - Analyze the query, previous reasoning steps, and observations. - Reflect on 5-7 different ways to solve the given query or task. Think carefully about each solution before picking the best one. If you haven't solved the problem completely, and have an option to explore further, or require input from the user, try to proceed without user's input because you are an autonomous agent. - Decide on the next action: use a tool or provide a final answer. {{end}} ## Command Structuring Guidelines: **IMPORTANT:** - When generating kubectl commands, ALWAYS place the verb (e.g., get, apply, delete) immediately after `kubectl`. - Example: - ✅ Correct: `kubectl get pods` - ✅ Correct: `kubectl get pods --all-namespaces` - ❌ Incorrect: `get pods` - ❌ Incorrect: `get pods --all-namespaces` - Do NOT place flags or options before the verb. - Example: - ✅ Correct: `kubectl get pods --namespace=default` - ❌ Incorrect: `kubectl --namespace=default get pods` - This ensures commands are properly recognized and filtered by the system. - Prefer the command that does not require any interactive input. {{if .SessionIsInteractive}} ## Resource Manifest Generation Guidelines: **CRITICAL**: NEVER generate or create Kubernetes manifests without FIRST gathering ALL required specifics from the user and cluster state. This is a MANDATORY step that cannot be skipped. ### MANDATORY Information Collection Process: Before creating ANY manifest, you MUST: 1. **Check Cluster State**: - Run `kubectl get namespaces` to show available namespaces - Run `kubectl get nodes` to understand cluster capacity - Run `kubectl get storageclass` if storage is involved - Check existing resources with relevant `kubectl get` commands 2. **Ask User for Missing Specifics** (DO NOT assume defaults): - **Namespace**: "Which namespace should I deploy this to?" (show available options) - **Container Images**: "Which specific image version should I use?" (e.g., postgres:14, postgres:15, postgres:latest) - **Storage Size**: "How much storage do you need?" (if persistent storage required) - **Resource Limits**: "What CPU/memory limits should I set?" - **Service Exposure**: "How should this be exposed?" (ClusterIP, NodePort, LoadBalancer) - **Environment Variables**: "Do you need any specific environment variables or configurations?" - **Security**: "Do you need specific passwords, secrets, or service accounts?" 3. **Present Summary for Confirmation**: After gathering details, present a summary like: ``` **Deployment Summary:** - Namespace: [specified namespace] - Image: [specific image:tag] - Storage: [size] with [storage class] - Resources: [CPU/memory limits] - Service: [exposure type] - Security: [password/secret configuration] Should I proceed with creating these resources? Please confirm. ``` ### STRICT Manifest Creation Rules: - **NEVER** generate manifests with assumed defaults without user confirmation - **NEVER** skip the information gathering phase - **NEVER** proceed without explicit user confirmation of the configuration - **ALWAYS** ask specific questions about unclear requirements - **ALWAYS** show available options (namespaces, storage classes, etc.) - **ALWAYS** confirm the final configuration before creating resources ### Required Information to Collect: 1. **Namespace**: Check existing namespaces and ask which namespace to use if not specified 2. **Container Images**: - Verify image availability and tags - Check for specific version requirements - Validate image registry accessibility 3. **Ports and Services**: - Identify required container ports - Determine service type (ClusterIP, NodePort, LoadBalancer) - Check for existing services that might conflict 4. **Resource Requirements**: - CPU and memory requests/limits - Storage requirements (PVCs, volumes) - Node selection criteria (selectors, affinity) 5. **Environment Configuration**: - Required environment variables - ConfigMaps and Secrets needed - Service accounts and RBAC requirements 6. **Dependencies**: - Check for existing resources that need to be referenced - Verify network policies don't block connections - Ensure required CRDs are installed {{end}} ## Remember: - Fetch current state of kubernetes resources relevant to user's query. - If using a kubectl command ensure that verb is always prefixed by `kubectl`. - Prefer the tool usage that does not require any interactive input. - For creating new resources, try to create the resource using the tools available. DO NOT ask the user to create the resource. - Use tools when you need more information. Do not respond with the instructions on how to use the tools or what commands to run, instead just use the tool. - Provide a final answer only when you're confident you have sufficient information. - Provide clear, concise, and accurate responses. - Feel free to respond with emojis where appropriate. ================================================ FILE: pkg/api/models.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package api import ( "fmt" "time" ) type Session struct { ID string Name string ProviderID string ModelID string Messages []*Message AgentState AgentState CreatedAt time.Time LastModified time.Time ChatMessageStore ChatMessageStore // MCP status information MCPStatus *MCPStatus } type AgentState string const ( AgentStateIdle AgentState = "idle" AgentStateWaitingForInput AgentState = "waiting-for-input" AgentStateRunning AgentState = "running" AgentStateInitializing AgentState = "initializing" AgentStateDone AgentState = "done" AgentStateExited AgentState = "exited" ) type MessageType string const ( MessageTypeText MessageType = "text" MessageTypeError MessageType = "error" MessageTypeToolCallRequest MessageType = "tool-call-request" MessageTypeToolCallResponse MessageType = "tool-call-response" MessageTypeUserInputRequest MessageType = "user-input-request" MessageTypeUserInputResponse MessageType = "user-input-response" MessageTypeUserChoiceRequest MessageType = "user-choice-request" MessageTypeUserChoiceResponse MessageType = "user-choice-response" MessageTypeSessionPickerRequest MessageType = "session-picker-request" MessageTypeSessionPickerResponse MessageType = "session-picker-response" ) type Message struct { ID string Source MessageSource Type MessageType Payload any Timestamp time.Time } type MessageSource string const ( MessageSourceUser MessageSource = "user" MessageSourceAgent MessageSource = "agent" MessageSourceModel MessageSource = "model" ) type UserChoiceRequest struct { Prompt string Options []UserChoiceOption } type UserChoiceOption struct { Label string `json:"label,omitempty"` Value string `json:"value,omitempty"` } type UserChoiceResponse struct { Choice int `json:"choice"` } type UserInputResponse struct { Query string `json:"query"` } // SessionPickerRequest is sent to show an interactive session picker type SessionPickerRequest struct { Sessions []SessionInfo `json:"sessions"` } // SessionInfo contains display information for a session type SessionInfo struct { ID string `json:"id"` Name string `json:"name,omitempty"` ModelID string `json:"modelId,omitempty"` ProviderID string `json:"providerId,omitempty"` CreatedAt time.Time `json:"createdAt"` LastModified time.Time `json:"lastModified"` MessageCount int `json:"messageCount"` } // SessionPickerResponse is sent when user selects a session type SessionPickerResponse struct { SessionID string `json:"sessionId"` Cancelled bool `json:"cancelled"` } // MCPStatus represents the overall status of MCP servers and tools type MCPStatus struct { ServerInfoList []ServerConnectionInfo `json:"serverInfoList,omitempty"` TotalServers int `json:"totalServers,omitempty"` ConnectedCount int `json:"connectedCount,omitempty"` FailedCount int `json:"failedCount,omitempty"` TotalTools int `json:"totalTools,omitempty"` ClientEnabled bool `json:"clientEnabled,omitempty"` } // ServerConnectionInfo holds connection status for a single MCP server type ServerConnectionInfo struct { Name string `json:"name,omitempty"` Command string `json:"command,omitempty"` IsLegacy bool `json:"isLegacy,omitempty"` IsConnected bool `json:"isConnected,omitempty"` AvailableTools []MCPTool `json:"availableTools,omitempty"` } // MCPTool represents an MCP tool with basic information type MCPTool struct { Name string `json:"name,omitempty"` Description string `json:"description,omitempty"` Server string `json:"server,omitempty"` } // ChatMessageStore defines the interface for managing storage of chat messages of a session. type ChatMessageStore interface { AddChatMessage(record *Message) error SetChatMessages(newHistory []*Message) error ChatMessages() []*Message ClearChatMessages() error } func (s *Session) AllMessages() []*Message { if s.ChatMessageStore == nil { return nil } return s.ChatMessageStore.ChatMessages() } func (s *Session) String() string { return fmt.Sprintf("Session ID: %s\nProvider: %s\nModel: %s\nCreated At: %s\nLast Modified: %s\nAgent State: %s", s.ID, s.ProviderID, s.ModelID, s.CreatedAt.Format(time.RFC3339), s.LastModified.Format(time.RFC3339), s.AgentState) } ================================================ FILE: pkg/journal/context.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package journal import ( "context" ) type contextKey string const RecorderKey contextKey = "journal-recorder" // RecorderFromContext extracts the recorder from the given context func RecorderFromContext(ctx context.Context) Recorder { recorder, ok := ctx.Value(RecorderKey).(Recorder) if !ok { return &LogRecorder{} } return recorder } // ContextWithRecorder adds the recorder to the given context func ContextWithRecorder(ctx context.Context, recorder Recorder) context.Context { return context.WithValue(ctx, RecorderKey, recorder) } ================================================ FILE: pkg/journal/loader.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package journal import ( "bufio" "bytes" "fmt" "io" "os" "sigs.k8s.io/yaml" ) // ParseEventsFromFile will read the events from the given file path func ParseEventsFromFile(p string) ([]*Event, error) { f, err := os.Open(p) if err != nil { return nil, fmt.Errorf("opening file %q: %w", p, err) } defer f.Close() return ParseEvents(f) } // ParseEvents will read the events from the reader func ParseEvents(r io.Reader) ([]*Event, error) { var events []*Event scanner := bufio.NewScanner(r) scanner.Split(splitYAML) for scanner.Scan() { b := scanner.Bytes() event := &Event{} if err := yaml.Unmarshal(b, &event); err != nil { return nil, fmt.Errorf("parsing yaml: %w", err) } if event != nil { events = append(events, event) } } return events, nil } var yamlSep = []byte("\n---\n") // splitYAML is a split function for a Scanner that returns each object in a yaml multi-object doc. func splitYAML(data []byte, atEOF bool) (advance int, token []byte, err error) { if atEOF && len(data) == 0 { return 0, nil, nil } if i := bytes.Index(data, yamlSep); i >= 0 { // We have a full object. return i + len(yamlSep), data[0:i], nil } // If we're at EOF, we have a final, non-terminated object. Return it. if atEOF { return len(data), data, nil } // Request more data. return 0, nil, nil } ================================================ FILE: pkg/journal/log.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package journal import ( "context" "k8s.io/klog/v2" ) type LogRecorder struct { } func (r *LogRecorder) Write(ctx context.Context, event *Event) error { log := klog.FromContext(ctx) log.V(2).Info("Tracing event", "event", event) return nil } func (r *LogRecorder) Close() error { return nil } ================================================ FILE: pkg/journal/recorder.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package journal import ( "bytes" "context" "fmt" "io" "os" "time" "sigs.k8s.io/yaml" ) // Recorder is an interface for recording a structured log of the agent's actions and observations. type Recorder interface { io.Closer // Write will add an event to the recorder. Write(ctx context.Context, event *Event) error } // FileRecorder writes a structured log of the agent's actions and observations to a file. type FileRecorder struct { f *os.File } // NewFileRecorder creates a new FileRecorder that writes to the given file. func NewFileRecorder(path string) (*FileRecorder, error) { file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) if err != nil { return nil, fmt.Errorf("opening file: %w", err) } return &FileRecorder{ f: file, }, nil } // Close closes the file. func (r *FileRecorder) Close() error { return r.f.Close() } func (r *FileRecorder) Write(ctx context.Context, event *Event) error { if event.Timestamp.IsZero() { event.Timestamp = time.Now() } yamlBytes, err := yaml.Marshal(event) if err != nil { return fmt.Errorf("marshalling event: %w", err) } var b bytes.Buffer b.Write(yamlBytes) b.Write([]byte("\n\n---\n\n")) _, err = r.f.Write(b.Bytes()) return err } type Event struct { Timestamp time.Time `json:"timestamp"` Action string `json:"action"` Payload any `json:"payload,omitempty"` } const ( ActionHTTPRequest = "http.request" ActionHTTPResponse = "http.response" ActionHTTPError = "http.error" ) // ActionUIRender is for an event that indicates we wrote output to the UI const ActionUIRender = "ui.render" // GetString is a helper to get a string value from the Payload func (e *Event) GetString(key string) (string, bool) { if e.Payload == nil { return "", false } m, ok := e.Payload.(map[string]any) if !ok { return "", false } v, ok := m[key] if !ok { return "", false } s, ok := v.(string) if !ok { return "", false } return s, true } ================================================ FILE: pkg/mcp/README.md ================================================ # MCP (Model Context Protocol) Client This package provides functionality to interact with MCP (Model Context Protocol) servers from `kubectl-ai`. ## Overview The MCP client allows `kubectl-ai` to connect to MCP servers, discover available tools, and execute them. This enables integration with various services and systems that expose their functionality through the MCP protocol. ## Features - Connect to multiple MCP servers simultaneously - Support for both local (stdio-based) and remote (HTTP-based) MCP servers - Authentication support for HTTP-based servers (Basic, Bearer Token, API Key) - Automatic discovery of available tools from connected servers - Execute tools on MCP servers with parameter conversion - Configuration-based server management - Generic parameter name and type conversion (snake_case → camelCase, intelligent type inference) - Synchronous initialization ensuring tools are available before conversation starts ## Configuration MCP server configurations are stored in `~/.config/kubectl-ai/mcp.yaml`. If this file doesn't exist, a default configuration will be created automatically. ### Default Configuration By default, the MCP client is configured with sequential thinking MCP server: ```yaml servers: - name: sequential-thinking command: npx args: - -y - "@modelcontextprotocol/server-sequential-thinking" ``` ### Configuration Format The configuration file uses YAML format and supports both local (stdio-based) and remote (HTTP-based) MCP servers: #### Local (stdio-based) Server Configuration ```yaml servers: - name: server-name command: path-to-server-binary args: - --flag1 - value1 env: ENV_VAR: value ``` #### Remote (HTTP-based) Server Configuration ```yaml servers: - name: remote-server url: "https://mcp-server.example.com/" timeout: 30 # Optional: Timeout in seconds use_streaming: true # Optional: Use streaming HTTP client # Optional authentication auth: type: "bearer" # Options: "basic", "bearer", "api-key" token: "${YOUR_ENV_VAR}" # Will be read from YOUR_ENV_VAR environment variable # Optional custom headers headers: X-Custom-Header: "custom-value" X-API-Version: "v1" ``` ### Authentication Options Remote MCP servers support different authentication methods: 1. **Bearer Token**: ```yaml auth: type: "bearer" token: "your-bearer-token" ``` 2. **Basic Authentication**: ```yaml auth: type: "basic" username: "username" password: "password" ``` 3. **API Key**: ```yaml auth: type: "api-key" api_key: "your-api-key" header_name: "X-Api-Key" # Optional: Defaults to X-Api-Key ``` ### Custom Headers Remote MCP servers support custom HTTP headers for additional configuration or authentication requirements: ```yaml servers: - name: remote-server url: "https://mcp-server.example.com/" headers: X-Custom-Header: "custom-value" X-API-Version: "v1" User-Agent: "kubectl-ai/1.0" Accept-Language: "en-US" ``` **Key Points:** - Custom headers are applied to all HTTP requests sent to the MCP server - Headers can be combined with authentication methods - Authentication headers (e.g., `Authorization`) may override custom headers if both are specified - All header values are strings - Headers are case-sensitive as per HTTP specification **Common Use Cases:** - API versioning headers (e.g., `X-API-Version`) - Custom user agent strings - Request tracking headers (e.g., `X-Request-ID`) - Content negotiation headers (e.g., `Accept`, `Accept-Language`) ### Environment Variable Support Sensitive information like tokens and passwords can be read from environment variables using the `${VAR_NAME}` syntax in the configuration file. You can also set environment variables with the prefix `MCP_SERVER_NAME_` to override configuration values. ## Usage Enable MCP client functionality with the `--mcp-client` flag: ```bash kubectl-ai --mcp-client ``` ### Checking Server Status When you run kubectl-ai with the MCP client enabled, you'll see information about connected servers: ```txt MCP Server Status: Successfully connected to 2 MCP server(s) (2 tools discovered) • sequential-thinking (npx) - Connected, Tools: sequentialthinking • fetch (remote) - Connected, Tools: fetch ``` MCP servers are automatically discovered and their tools made available to the AI. The system handles: - **Parameter conversion**: Automatically converts snake_case parameters to camelCase - **Type inference**: Intelligently converts string parameters to numbers/booleans based on naming patterns - **Error handling**: Graceful fallbacks for connection issues ### Custom Server Examples To add custom MCP servers, edit the configuration file at `~/.config/kubectl-ai/mcp.yaml`: You can combine both local and remote servers in your configuration: ```yaml servers: - name: sequential-thinking command: npx args: - -y - '@modelcontextprotocol/server-sequential-thinking' - name: cloudflare-documentation url: https://docs.mcp.cloudflare.com/mcp ``` ### Environment Variables You can configure the following environment variables to customize MCP client behavior: - `KUBECTL_AI_MCP_CONFIG`: Override the default configuration file path - `MCP__`: Set environment variables for specific servers ## Parameter Conversion The MCP client automatically handles parameter name and type conversion to ensure compatibility with different MCP servers: ### Name Conversion - Converts snake_case parameter names to camelCase - Example: `thought_number` → `thoughtNumber` ### Type Conversion Parameters are intelligently converted based on naming patterns: **Numbers:** Parameters containing `number`, `count`, `total`, `max`, `min`, `limit` **Booleans:** Parameters starting with `is`, `has`, `needs`, `enable` or containing `required`, `enabled` ### Fallback Behavior - If type conversion fails, the original value is preserved - Unknown servers use generic conversion rules - No configuration required - works automatically with any MCP server ## Implementation Details ### Client The `Client` struct represents a connection to an MCP server. It provides methods to: - Connect to the server - List available tools - Execute tools - Close the connection ### Manager The `Manager` struct manages multiple MCP client connections. It provides: - Connection management for multiple servers - Tool discovery across all connected servers - Thread-safe operations ### Configuration The `Config` struct handles loading and saving MCP server configurations from disk. The configuration is automatically loaded from `~/.config/kubectl-ai/mcp.yaml` when needed. ## Integration with kubectl-ai The MCP client is integrated with `kubectl-ai` to automatically discover and use tools from configured MCP servers. The system: 1. **Loads configuration** from `~/.config/kubectl-ai/mcp.yaml` on startup 2. **Connects synchronously** to all configured MCP servers (when `--mcp-client` flag is used) 3. **Registers tools** before the conversation starts, ensuring they're immediately available 4. **Converts parameters** automatically using generic snake_case → camelCase conversion 5. **Handles execution** with proper error handling and result formatting 6. **Displays status** showing connected servers and available tool counts 📖 **For practical multi-server orchestration examples and security automation workflows, see the [MCP Client Integration Guide](../../docs/mcp-client.md).** ## Security Considerations - MCP servers can execute arbitrary commands with the same permissions as the `kubectl-ai` process - Only connect to trusted MCP servers - The configuration file has strict permissions (0600) by default - Be cautious when adding environment variables with sensitive information ## Troubleshooting ### Common Issues **MCP tools are not available:** - Ensure you're using the `--mcp-client` flag - Check that `~/.config/kubectl-ai/mcp.yaml` exists and is valid (created by default) - Verify MCP servers are installed (e.g., `npx` commands work) **Connection failures:** - Check network connectivity - Ensure server commands and paths are correct in configuration - Verify environment variables are properly set **Parameter conversion issues:** - The system automatically converts snake_case → camelCase - String parameters are converted to numbers/booleans based on naming patterns - Fallback behavior preserves original values if conversion fails ### Debug Information - Use `-v=1` for basic MCP operation logging - Use `-v=2` for detailed connection and tool discovery info - Check server status in the startup message - Tool counts are displayed for each connected server ================================================ FILE: pkg/mcp/client.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "context" "fmt" "reflect" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" mcpclient "github.com/mark3labs/mcp-go/client" mcp "github.com/mark3labs/mcp-go/mcp" "k8s.io/klog/v2" ) // =================================================================== // Client Types and Factory Functions // =================================================================== // Client represents an MCP client that can connect to MCP servers. // It is a wrapper around the MCPClient interface for backward compatibility. type Client struct { // Name is a friendly name for this MCP server connection Name string // The actual client implementation (stdio or HTTP) impl MCPClient // client is the underlying MCP library client client *mcpclient.Client } // Tool represents an MCP tool with optional server information. type Tool struct { Name string `json:"name"` Description string `json:"description,omitempty"` Server string `json:"server,omitempty"` InputSchema *gollm.Schema `json:"inputSchema,omitempty"` } // NewClient creates a new MCP client with the given configuration. // This function supports both stdio and HTTP-based MCP servers. func NewClient(config ClientConfig) *Client { // Create the appropriate implementation based on configuration var impl MCPClient if config.URL != "" { // HTTP-based client impl = NewHTTPClient(config) } else { // Stdio-based client impl = NewStdioClient(config) } return &Client{ Name: config.Name, impl: impl, } } // CreateStdioClient creates a new stdio-based MCP client (for backward compatibility). func CreateStdioClient(name, command string, args []string, env map[string]string) *Client { // Convert env map to slice of KEY=value strings var envSlice []string for k, v := range env { envSlice = append(envSlice, fmt.Sprintf("%s=%s", k, v)) } config := ClientConfig{ Name: name, Command: command, Args: args, Env: envSlice, } return NewClient(config) } // =================================================================== // Main Client Interface Methods // =================================================================== // Connect establishes a connection to the MCP server. // This delegates to the appropriate implementation (stdio or HTTP). func (c *Client) Connect(ctx context.Context) error { klog.V(2).InfoS("Connecting to MCP server", "name", c.Name) // Delegate to the implementation if err := c.impl.Connect(ctx); err != nil { return err } // Store the underlying client for backward compatibility c.client = c.impl.getUnderlyingClient() klog.V(2).InfoS("Successfully connected to MCP server", "name", c.Name) return nil } // Close closes the connection to the MCP server. func (c *Client) Close() error { if c.impl == nil { return nil // Not initialized } klog.V(2).InfoS("Closing connection to MCP server", "name", c.Name) // Delegate to implementation err := c.impl.Close() c.client = nil // Clear reference to underlying client if err != nil { return fmt.Errorf("closing MCP client: %w", err) } return nil } // ListTools lists all available tools from the MCP server. func (c *Client) ListTools(ctx context.Context) ([]Tool, error) { if err := c.ensureConnected(); err != nil { return nil, err } // Delegate to implementation tools, err := c.impl.ListTools(ctx) if err != nil { return nil, err } klog.V(2).InfoS("Listed tools from MCP server", "count", len(tools), "server", c.Name) return tools, nil } // CallTool calls a tool on the MCP server and returns the result as a string. // The arguments should be a map of parameter names to values that will be passed to the tool. func (c *Client) CallTool(ctx context.Context, toolName string, arguments map[string]interface{}) (string, error) { klog.V(2).InfoS("Calling MCP tool", "server", c.Name, "tool", toolName, "args", arguments) if err := c.ensureConnected(); err != nil { return "", err } // Delegate to implementation return c.impl.CallTool(ctx, toolName, arguments) } // =================================================================== // Tool Factory Functions and Methods // =================================================================== // WithServer returns a copy of the tool with server information added. func (t Tool) WithServer(server string) Tool { copy := t copy.Server = server return copy } // ID returns a unique identifier for the tool. func (t Tool) ID() string { if t.Server != "" { return fmt.Sprintf("%s@%s", t.Name, t.Server) } return t.Name } // String returns a human-readable representation of the tool. func (t Tool) String() string { if t.Server != "" { return fmt.Sprintf("%s (from %s)", t.Name, t.Server) } return t.Name } // AsBasicTool returns the tool without server information (for client.ListTools compatibility). func (t Tool) AsBasicTool() Tool { copy := t copy.Server = "" return copy } // IsFromServer checks if the tool belongs to a specific server. func (t Tool) IsFromServer(server string) bool { return t.Server == server } // convertMCPToolsToTools converts MCP library tools to our Tool type. func convertMCPToolsToTools(mcpTools []mcp.Tool) ([]Tool, error) { tools := make([]Tool, 0, len(mcpTools)) for _, mcpTool := range mcpTools { tool := Tool{ Name: mcpTool.Name, Description: mcpTool.Description, } // TODO: Annotations (give hints about e.g. read-only, destructive, idempotent, open-world) if mcpTool.InputSchema.Type != "" { schema, err := convertMCPInputSchema(&mcpTool.InputSchema) if err != nil { return nil, fmt.Errorf("converting MCP input schema to tool input schema: %w", err) } tool.InputSchema = schema } else { // TODO: Use RawInputSchema if available // klog.Warningf("no input schema for tool %s", mcpTool.Name) return nil, fmt.Errorf("no input schema for tool %s", mcpTool.Name) } tools = append(tools, tool) } return tools, nil } func convertMCPInputSchema(mcpInputSchema *mcp.ToolInputSchema) (*gollm.Schema, error) { gollmSchema := &gollm.Schema{} switch mcpInputSchema.Type { case "string": gollmSchema.Type = gollm.TypeString // case "number": // gollmSchema.Type = gollm.TypeNumber case "boolean": gollmSchema.Type = gollm.TypeBoolean case "object": gollmSchema.Type = gollm.TypeObject default: return nil, fmt.Errorf("unexpected MCP input schema type: %s", mcpInputSchema.Type) } if mcpInputSchema.Properties != nil { gollmSchema.Properties = make(map[string]*gollm.Schema) for key, value := range mcpInputSchema.Properties { if valueSchema, ok := value.(mcp.ToolInputSchema); ok { gollmValue, err := convertMCPInputSchema(&valueSchema) if err != nil { return nil, fmt.Errorf("converting MCP input schema to tool input schema: %w", err) } gollmSchema.Properties[key] = gollmValue } else if valueMap, ok := value.(map[string]interface{}); ok && valueMap != nil { gollmValue, err := convertMCPMapSchema(key, valueMap) if err != nil { return nil, fmt.Errorf("converting MCP input schema to tool input schema: %w", err) } gollmSchema.Properties[key] = gollmValue } else { return nil, fmt.Errorf("unexpected input schema type for %q: %T %+v", key, value, value) } } } gollmSchema.Required = mcpInputSchema.Required return gollmSchema, nil } func convertMCPMapSchema(key string, schemaMap map[string]interface{}) (*gollm.Schema, error) { if schemaMap == nil { return nil, fmt.Errorf("schema map is nil for key %q", key) } gollmSchema := &gollm.Schema{} if descriptionObj, ok := schemaMap["description"]; ok { description, ok := descriptionObj.(string) if !ok { return nil, fmt.Errorf("unexpected description for key %q: %+v", key, schemaMap) } gollmSchema.Description = description } mcpType, ok := schemaMap["type"].(string) if !ok { // Fallback: treat any unrecognized schema as generic object klog.V(2).InfoS("Unrecognized schema format, treating as object", "key", key) gollmSchema.Type = gollm.TypeObject return gollmSchema, nil } switch mcpType { case "string": gollmSchema.Type = gollm.TypeString case "number": gollmSchema.Type = gollm.TypeNumber case "integer": gollmSchema.Type = gollm.TypeNumber case "boolean": gollmSchema.Type = gollm.TypeBoolean case "array": items, ok := schemaMap["items"].(map[string]interface{}) if !ok { return nil, fmt.Errorf("did not find items for array: key %q: %+v", key, schemaMap) } itemsSchema, err := convertMCPMapSchema(key+".items", items) if err != nil { return nil, fmt.Errorf("converting MCP input schema to tool input schema: %w", err) } gollmSchema.Type = gollm.TypeArray gollmSchema.Items = itemsSchema case "object": gollmSchema.Type = gollm.TypeObject gollmSchema.Properties = make(map[string]*gollm.Schema) if propertiesObj, ok := schemaMap["properties"]; ok && propertiesObj != nil { properties, ok := propertiesObj.(map[string]interface{}) if !ok { return nil, fmt.Errorf("properties field is not a map for key %q: %+v", key, schemaMap) } for key, value := range properties { valueMap, ok := value.(map[string]interface{}) if !ok { return nil, fmt.Errorf("property value is not a map for key %q: %+v", key, value) } propertySchema, err := convertMCPMapSchema(key, valueMap) if err != nil { return nil, fmt.Errorf("converting MCP input schema to tool input schema: %w", err) } gollmSchema.Properties[key] = propertySchema } } default: return nil, fmt.Errorf("unexpected input schema type %q for key %q: %+v", mcpType, key, schemaMap) } return gollmSchema, nil } // =================================================================== // Common Functions // =================================================================== // ensureClientConnected checks if the client is connected. func ensureClientConnected(client *mcpclient.Client) error { if client == nil { return fmt.Errorf("client not connected") } return nil } // initializeClientConnection initializes the MCP connection with proper handshake. func initializeClientConnection(ctx context.Context, client *mcpclient.Client) error { initCtx, cancel := context.WithTimeout(ctx, DefaultConnectionTimeout) defer cancel() // Create initialize request with the structure expected by v0.31.0 initReq := mcp.InitializeRequest{ // The structure might differ in v0.31.0 - adapt as needed // This is a placeholder that will be updated when the actual API is known } _, err := client.Initialize(initCtx, initReq) if err != nil { return fmt.Errorf("initializing MCP client: %w", err) } return nil } // verifyClientConnection verifies the connection works by testing tool listing. func verifyClientConnection(ctx context.Context, client *mcpclient.Client) error { verifyCtx, cancel := context.WithTimeout(ctx, DefaultConnectionTimeout) defer cancel() // Try to list tools as a basic connectivity test _, err := client.ListTools(verifyCtx, mcp.ListToolsRequest{}) if err != nil { return fmt.Errorf("listing tools: %w", err) } return nil } // cleanupClient closes the client connection safely. func cleanupClient(client **mcpclient.Client) { if *client != nil { _ = (*client).Close() // Ignore errors on cleanup *client = nil } } // processToolResponse processes a tool call response and extracts the text result. // This function works with any MCP response object that has the expected fields. func processToolResponse(result any) (string, error) { // Use reflection to safely access fields rv := reflect.ValueOf(result) // Handle pointer to struct if rv.Kind() == reflect.Ptr { rv = rv.Elem() } if rv.Kind() != reflect.Struct { return "", fmt.Errorf("unexpected response type: %T", result) } // Check for IsError field isErrorField := rv.FieldByName("IsError") if isErrorField.IsValid() && isErrorField.Kind() == reflect.Bool { isError := isErrorField.Bool() // Handle error response if isError { // Extract error message errorMsg := fmt.Sprintf("%+v", result) // Try to get message from Content field contentField := rv.FieldByName("Content") if contentField.IsValid() && contentField.Len() > 0 { if content := contentField.Index(0).Interface(); content != nil { if textContent, ok := mcp.AsTextContent(content); ok { errorMsg = textContent.Text } } } // Return JSON error data instead of Go error return fmt.Sprintf(`{"error": true, "message": %q, "status": "failed"}`, errorMsg), nil } } // Check for Content field contentField := rv.FieldByName("Content") if contentField.IsValid() && contentField.Len() > 0 { // Let's rely on the AsTextContent method from MCP package // which handles the specific response format content := contentField.Index(0).Interface() if textContent, ok := mcp.AsTextContent(content); ok { return textContent.Text, nil } } // If we couldn't extract text content, return a generic message return "Tool executed successfully, but no text content was returned", nil } // listClientTools implements the common ListTools functionality shared by both client types. func listClientTools(ctx context.Context, client *mcpclient.Client, serverName string) ([]Tool, error) { if err := ensureClientConnected(client); err != nil { return nil, err } // Call the ListTools method on the MCP server result, err := client.ListTools(ctx, mcp.ListToolsRequest{}) if err != nil { return nil, fmt.Errorf("listing tools: %w", err) } // Convert the result using the helper function tools, err := convertMCPToolsToTools(result.Tools) if err != nil { return nil, fmt.Errorf("parsing tools from MCP server: %w", err) } // Add the server name to each tool for i := range tools { tools[i].Server = serverName } return tools, nil } ================================================ FILE: pkg/mcp/config.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "fmt" "os" "path/filepath" "runtime" "strings" "k8s.io/klog/v2" "sigs.k8s.io/yaml" ) // Config represents the complete MCP client configuration file type Config struct { // Servers is a list of MCP server configurations Servers []ServerConfig `yaml:"servers,omitempty"` } // ServerConfig represents the configuration for a single MCP server type ServerConfig struct { // Name is a friendly name for this MCP server Name string `yaml:"name"` // Command is the command to execute for stdio-based MCP servers Command string `yaml:"command"` // Args are the arguments to pass to the command Args []string `yaml:"args,omitempty"` // Env are the environment variables to set for the command Env map[string]string `yaml:"env,omitempty"` // URL is the URL for HTTP-based MCP servers URL string `yaml:"url,omitempty"` // Auth is the authentication configuration for HTTP-based MCP servers Auth *AuthConfig `yaml:"auth,omitempty"` // OAuthConfig is the OAuth configuration for HTTP-based MCP servers OAuthConfig *OAuthConfig `yaml:"oauth,omitempty"` // Timeout is the timeout in seconds for HTTP requests Timeout int `yaml:"timeout,omitempty"` // UseStreaming enables streaming HTTP for better performance UseStreaming bool `yaml:"use_streaming,omitempty"` // SkipVerify skips TLS certificate verification for HTTPS connections SkipVerify bool `yaml:"skip_verify,omitempty"` } // =================================================================== // Configuration loading and management functions // =================================================================== // loadDefaultConfig loads the default configuration from the embedded file func loadDefaultConfig() (*Config, error) { // This path is relative to the module root defaultConfigPath := filepath.Join("pkg", "mcp", "default_config.yaml") // Read the file data, err := os.ReadFile(defaultConfigPath) if err != nil { return nil, fmt.Errorf("reading default config file: %w", err) } var config Config if err := yaml.Unmarshal(data, &config); err != nil { return nil, fmt.Errorf("parsing default config: %w", err) } return &config, nil } // DefaultConfigPath returns the default path to the MCP config file func DefaultConfigPath() (string, error) { // Get the home directory first home, err := os.UserHomeDir() if err != nil { return "", fmt.Errorf("getting user home directory: %w", err) } var configPath string // Handle different operating systems switch runtime.GOOS { case "windows": // On Windows, use %APPDATA%\kubectl-ai\mcp.yaml appData := os.Getenv("APPDATA") if appData == "" { appData = filepath.Join(home, "AppData", "Roaming") } configPath = filepath.Join(appData, "kubectl-ai", "mcp.yaml") default: // On Unix-like systems, use XDG_CONFIG_HOME/kubectl-ai/mcp.yaml configDir := os.Getenv("XDG_CONFIG_HOME") if configDir == "" { configDir = filepath.Join(home, ".config") } configPath = filepath.Join(configDir, "kubectl-ai", "mcp.yaml") } return configPath, nil } // LoadConfig loads the MCP configuration from the given path and applies environment variable overrides // If path is empty, the default config path is used // If the file doesn't exist, it creates a default configuration file func LoadConfig(path string) (*Config, error) { if path == "" { var err error path, err = DefaultConfigPath() if err != nil { return nil, err } } // If the file doesn't exist, create it with default configuration if _, err := os.Stat(path); os.IsNotExist(err) { // Create the directory if it doesn't exist dir := filepath.Dir(path) if err := os.MkdirAll(dir, ConfigDirPermissions); err != nil { return nil, fmt.Errorf("creating config directory: %w", err) } // Read the default config from the embedded file defaultConfig, err := loadDefaultConfig() if err != nil { return nil, fmt.Errorf("loading default config: %w", err) } // Save it to the config path if err := defaultConfig.Save(path); err != nil { return nil, fmt.Errorf("saving default config: %w", err) } // Apply environment variable overrides applyEnvironmentVariables(defaultConfig) return defaultConfig, nil } // Read the file data, err := os.ReadFile(path) if err != nil { return nil, fmt.Errorf("reading config file: %w", err) } // Parse the YAML var config Config if err := yaml.Unmarshal(data, &config); err != nil { return nil, fmt.Errorf("parsing config file: %w", err) } // Validate the configuration if err := config.ValidateConfig(); err != nil { return nil, fmt.Errorf("invalid configuration: %w", err) } // Apply environment variable overrides applyEnvironmentVariables(&config) return &config, nil } // Save saves the configuration to the given path using atomic write func (c *Config) Save(path string) error { if path == "" { var err error path, err = DefaultConfigPath() if err != nil { return err } } // Ensure directory exists if err := os.MkdirAll(filepath.Dir(path), ConfigDirPermissions); err != nil { return fmt.Errorf("creating config directory: %w", err) } // Marshal the config to YAML data, err := yaml.Marshal(c) if err != nil { return fmt.Errorf("marshaling config: %w", err) } // Perform atomic write if err := atomicWriteFile(path, data, ConfigFilePermissions); err != nil { return fmt.Errorf("writing config file: %w", err) } klog.V(2).Info("Saved MCP configuration", "path", path) return nil } // atomicWriteFile writes data to a file atomically using a temporary file func atomicWriteFile(path string, data []byte, perm os.FileMode) error { dir := filepath.Dir(path) // Create temporary file in the same directory tmpFile, err := os.CreateTemp(dir, ".mcp-config-*") if err != nil { return fmt.Errorf("creating temporary file: %w", err) } tmpPath := tmpFile.Name() defer os.Remove(tmpPath) // Clean up on error // Write data to temporary file if _, err := tmpFile.Write(data); err != nil { tmpFile.Close() return fmt.Errorf("writing to temporary file: %w", err) } // Sync and close if err := tmpFile.Sync(); err != nil { tmpFile.Close() return fmt.Errorf("syncing temporary file: %w", err) } if err := tmpFile.Close(); err != nil { return fmt.Errorf("closing temporary file: %w", err) } // Set permissions if err := os.Chmod(tmpPath, perm); err != nil { return fmt.Errorf("setting file permissions: %w", err) } // Atomic rename return os.Rename(tmpPath, path) } // =================================================================== // Configuration validation functions // =================================================================== // ValidateConfig validates the entire configuration func (c *Config) ValidateConfig() error { if len(c.Servers) == 0 { return fmt.Errorf("no servers configured") } // Check for duplicate server names serverNames := make(map[string]bool) for i, server := range c.Servers { if err := ValidateServerConfig(server); err != nil { return fmt.Errorf("server %d (%s): %w", i, server.Name, err) } if serverNames[server.Name] { return fmt.Errorf("duplicate server name: %s", server.Name) } serverNames[server.Name] = true } return nil } // ValidateServerConfig validates a single server configuration func ValidateServerConfig(config ServerConfig) error { if config.Name == "" { return fmt.Errorf("server name cannot be empty") } // URL-based server (HTTP) or Command-based server (stdio) if config.URL == "" && config.Command == "" { return fmt.Errorf("either URL or Command must be specified") } // Additional validation could be added here: // - Check if command exists and is executable // - Validate environment variable format // - Check argument validity // - Validate URL format return nil } // =================================================================== // Environment variable handling functions // =================================================================== // applyEnvironmentVariables overrides config with environment variables func applyEnvironmentVariables(config *Config) { // Apply MCP server-specific environment variables for i := range config.Servers { applyServerEnvironment(&config.Servers[i]) } } // applyServerEnvironment applies environment variables for a specific MCP server func applyServerEnvironment(server *ServerConfig) { prefix := EnvMCPServerPrefix + strings.ToUpper(server.Name) + "_" // Process URL for HTTP servers if url := os.Getenv(prefix + "URL"); url != "" { server.URL = url klog.V(2).InfoS("Using URL from environment", "server", server.Name, "url", url) } // Process authentication for HTTP servers if server.URL != "" && server.Auth != nil { applyAuthEnvironmentVariables(server, prefix) } // Process command and arguments for stdio servers if server.Command != "" { applyCommandEnvironmentVariables(server, prefix) } } // applyAuthEnvironmentVariables applies authentication-related environment variables func applyAuthEnvironmentVariables(server *ServerConfig, prefix string) { // Process token for bearer auth if server.Auth.Type == "bearer" { if token := os.Getenv(prefix + "TOKEN"); token != "" { server.Auth.Token = token klog.V(2).InfoS("Using bearer token from environment", "server", server.Name) } } // Process API key for API key auth if server.Auth.Type == "api-key" { if apiKey := os.Getenv(prefix + "API_KEY"); apiKey != "" { server.Auth.ApiKey = apiKey klog.V(2).InfoS("Using API key from environment", "server", server.Name) } } // Process basic auth credentials if server.Auth.Type == "basic" { if username := os.Getenv(prefix + "USERNAME"); username != "" { server.Auth.Username = username } if password := os.Getenv(prefix + "PASSWORD"); password != "" { server.Auth.Password = password } } } // applyCommandEnvironmentVariables applies command-related environment variables func applyCommandEnvironmentVariables(server *ServerConfig, prefix string) { // Override command if cmd := os.Getenv(prefix + "COMMAND"); cmd != "" { server.Command = cmd klog.V(2).InfoS("Using command from environment", "server", server.Name, "command", cmd) } // Process environment variables for the server for _, env := range os.Environ() { if strings.HasPrefix(env, prefix) { parts := strings.SplitN(env, "=", 2) if len(parts) == 2 { varName := strings.TrimPrefix(parts[0], prefix) // Skip special variables that we process elsewhere if varName != "COMMAND" && varName != "URL" && varName != "TOKEN" && varName != "API_KEY" && varName != "USERNAME" && varName != "PASSWORD" { if server.Env == nil { server.Env = make(map[string]string) } server.Env[varName] = parts[1] klog.V(3).InfoS("Added environment variable from environment", "server", server.Name, "var", varName) } } } } } ================================================ FILE: pkg/mcp/constants.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import "time" // Timeout constants for MCP operations const ( // DefaultConnectionTimeout is the timeout for establishing connections to MCP servers DefaultConnectionTimeout = 30 * time.Second // DefaultVerificationTimeout is the timeout for verifying server connections DefaultVerificationTimeout = 10 * time.Second // DefaultPingTimeout is the timeout for ping operations DefaultPingTimeout = 5 * time.Second // DefaultStabilizationDelay is the delay to allow servers to stabilize after connection DefaultStabilizationDelay = 2 * time.Second ) // Error message templates const ( ErrServerConnectionFmt = "connecting to MCP server %q: %w" ErrServerCloseFmt = "closing MCP client %q: %w" ErrToolCallFmt = "calling tool %q: %w" ErrPathCheckFmt = "checking path %q: %w" ) // Client constants const ( ClientName = "kubectl-ai-mcp-client" ClientVersion = "1.0.0" ) // File permissions const ( ConfigFilePermissions = 0600 ConfigDirPermissions = 0755 ) // Constants for environment variables const ( // EnvMCPServerPrefix is the prefix for MCP server environment variables EnvMCPServerPrefix = "MCP_" ) ================================================ FILE: pkg/mcp/default_config.yaml ================================================ servers: - name: sequential-thinking command: npx args: - -y - "@modelcontextprotocol/server-sequential-thinking" ================================================ FILE: pkg/mcp/http_client.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "context" "crypto/tls" "encoding/base64" "fmt" "net/http" "time" mcpclient "github.com/mark3labs/mcp-go/client" "github.com/mark3labs/mcp-go/client/transport" mcp "github.com/mark3labs/mcp-go/mcp" "k8s.io/klog/v2" ) // =================================================================== // HTTP Client Implementation // =================================================================== // httpClient is an MCP client that communicates with HTTP-based MCP servers type httpClient struct { name string url string auth *AuthConfig oauthConfig *OAuthConfig timeout int useStreaming bool skipVerify bool headers map[string]string client *mcpclient.Client } // NewHTTPClient creates a new HTTP-based MCP client func NewHTTPClient(config ClientConfig) MCPClient { return &httpClient{ name: config.Name, url: config.URL, auth: config.Auth, oauthConfig: config.OAuthConfig, timeout: config.Timeout, useStreaming: config.UseStreaming, skipVerify: config.SkipVerify, headers: config.Headers, } } // getUnderlyingClient returns the underlying MCP client. func (c *httpClient) getUnderlyingClient() *mcpclient.Client { return c.client } // ensureConnected makes sure the client is connected. func (c *httpClient) ensureConnected() error { return ensureClientConnected(c.client) } // Name returns the name of this client. func (c *httpClient) Name() string { return c.name } // Connect establishes a connection to the HTTP MCP server. func (c *httpClient) Connect(ctx context.Context) error { klog.V(2).InfoS("Connecting to HTTP MCP server", "name", c.name, "url", c.url) if c.client != nil { return nil // Already connected } var client *mcpclient.Client var err error // Create the appropriate client based on configuration if c.oauthConfig != nil { client, err = c.createOAuthClient(ctx) } else if c.useStreaming { client, err = c.createStreamingClient() } else { client, err = c.createStandardClient() } if err != nil { return fmt.Errorf("creating HTTP MCP client: %w", err) } c.client = client // Initialize the connection if err := c.initializeConnection(ctx); err != nil { c.cleanup() return fmt.Errorf("initializing connection: %w", err) } // Verify connection if err := c.verifyConnection(ctx); err != nil { c.cleanup() return fmt.Errorf("verifying connection: %w", err) } klog.V(2).InfoS("Successfully connected to HTTP MCP server", "name", c.name) return nil } // createStreamingClient creates a streamable HTTP client for better performance func (c *httpClient) createStreamingClient() (*mcpclient.Client, error) { // Set up options for the HTTP client var options []transport.StreamableHTTPCOption // Add timeout if specified (only when not using custom client) if c.timeout > 0 { options = append(options, transport.WithHTTPTimeout(time.Duration(c.timeout)*time.Second)) } klog.V(2).InfoS("WARNING: TLS certificate verification is disabled", "server", c.name) // Handle TLS verification skip by creating custom HTTP client if c.skipVerify { klog.V(2).InfoS("WARNING: TLS certificate verification is disabled", "server", c.name) // Create custom HTTP client with TLS verification disabled customClient := &http.Client{ Transport: &http.Transport{ TLSClientConfig: &tls.Config{ InsecureSkipVerify: true, }, }, } // Add timeout to custom client if specified if c.timeout > 0 { customClient.Timeout = time.Duration(c.timeout) * time.Second } // Use the custom HTTP client options = append(options, transport.WithHTTPBasicClient(customClient)) } // Prepare headers map for authentication and custom headers headers := make(map[string]string) // Add custom headers from configuration first for key, value := range c.headers { headers[key] = value klog.V(3).InfoS("Using custom header for HTTP client", "server", c.name, "header", key) } // Add authentication headers if specified (may override custom headers) if c.auth != nil { switch c.auth.Type { case "basic": auth := "Basic " + base64.StdEncoding.EncodeToString([]byte(c.auth.Username+":"+c.auth.Password)) headers["Authorization"] = auth klog.V(3).InfoS("Using basic auth for HTTP client", "server", c.name) case "bearer": headers["Authorization"] = "Bearer " + c.auth.Token klog.V(3).InfoS("Using bearer auth for HTTP client", "server", c.name) case "api-key": headerName := "X-Api-Key" if c.auth.HeaderName != "" { headerName = c.auth.HeaderName } headers[headerName] = c.auth.ApiKey klog.V(3).InfoS("Using API key auth for HTTP client", "server", c.name) } } // Add headers if any were set if len(headers) > 0 { options = append(options, transport.WithHTTPHeaders(headers)) } klog.V(4).InfoS("Creating streamable HTTP client", "server", c.name, "url", c.url) client, err := mcpclient.NewStreamableHttpClient(c.url, options...) if err != nil { return nil, fmt.Errorf("creating streamable HTTP client: %w", err) } return client, nil } // createStandardClient creates a standard HTTP client func (c *httpClient) createStandardClient() (*mcpclient.Client, error) { // Standard client delegates to streaming client implementation for now // In the future, they might have different configurations return c.createStreamingClient() } // createOAuthClient creates an HTTP client with OAuth authentication func (c *httpClient) createOAuthClient(ctx context.Context) (*mcpclient.Client, error) { if c.oauthConfig == nil { return nil, fmt.Errorf("OAuth config required but not provided") } klog.V(3).InfoS("Creating OAuth HTTP client", "server", c.name, "client_id", c.oauthConfig.ClientID) // Set up options for the HTTP client var options []transport.StreamableHTTPCOption // Create OAuth configuration for the transport oauthCfg := transport.OAuthConfig{ ClientID: c.oauthConfig.ClientID, ClientSecret: c.oauthConfig.ClientSecret, Scopes: c.oauthConfig.Scopes, RedirectURI: c.oauthConfig.RedirectURL, // Use the token URL as the auth server metadata URL if available AuthServerMetadataURL: c.oauthConfig.TokenURL, } // Add OAuth configuration options = append(options, transport.WithHTTPOAuth(oauthCfg)) // Add timeout if specified if c.timeout > 0 { options = append(options, transport.WithHTTPTimeout(time.Duration(c.timeout)*time.Second)) } klog.V(4).InfoS("Creating OAuth streamable HTTP client", "server", c.name, "url", c.url) client, err := mcpclient.NewStreamableHttpClient(c.url, options...) if err != nil { return nil, fmt.Errorf("creating OAuth HTTP client: %w", err) } return client, nil } // initializeConnection initializes the MCP connection with proper handshake func (c *httpClient) initializeConnection(ctx context.Context) error { return initializeClientConnection(ctx, c.client) } // verifyConnection verifies the connection works by testing tool listing func (c *httpClient) verifyConnection(ctx context.Context) error { return verifyClientConnection(ctx, c.client) } // cleanup closes the client connection and resets the client state func (c *httpClient) cleanup() { cleanupClient(&c.client) } // Close closes the connection to the MCP server func (c *httpClient) Close() error { if c.client == nil { return nil // Already closed } klog.V(2).InfoS("Closing connection to HTTP MCP server", "name", c.name) err := c.client.Close() c.client = nil if err != nil { return fmt.Errorf("closing MCP client: %w", err) } return nil } // ListTools lists all available tools from the MCP server func (c *httpClient) ListTools(ctx context.Context) ([]Tool, error) { tools, err := listClientTools(ctx, c.client, c.name) if err != nil { return nil, err } klog.V(2).InfoS("Listed tools from HTTP MCP server", "count", len(tools), "server", c.name) return tools, nil } // CallTool calls a tool on the MCP server and returns the result as a string func (c *httpClient) CallTool(ctx context.Context, toolName string, arguments map[string]interface{}) (string, error) { klog.V(2).InfoS("Calling MCP tool via HTTP", "server", c.name, "tool", toolName) if err := c.ensureConnected(); err != nil { return "", err } // Create v0.31.0 compatible request request := mcp.CallToolRequest{ Params: mcp.CallToolParams{ Name: toolName, Arguments: arguments, }, } // Call the tool on the MCP server result, err := c.client.CallTool(ctx, request) if err != nil { return "", fmt.Errorf("error calling tool %s: %w", toolName, err) } return processToolResponse(result) } ================================================ FILE: pkg/mcp/interfaces.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "context" "fmt" mcpclient "github.com/mark3labs/mcp-go/client" ) // MCPClient defines the common interface for all MCP client implementations type MCPClient interface { // Name returns the name of this client Name() string // Connect establishes a connection to the MCP server Connect(ctx context.Context) error // Close closes the connection to the MCP server Close() error // ListTools lists all available tools from the MCP server ListTools(ctx context.Context) ([]Tool, error) // CallTool calls a tool on the MCP server and returns the result as a string CallTool(ctx context.Context, toolName string, arguments map[string]interface{}) (string, error) // ensureConnected makes sure the client is connected ensureConnected() error // getUnderlyingClient returns the underlying mcpclient.Client getUnderlyingClient() *mcpclient.Client } // ClientConfig contains all configuration options for MCP clients type ClientConfig struct { // Common fields Name string // For stdio-based clients Command string Args []string Env []string // For HTTP-based clients URL string Auth *AuthConfig OAuthConfig *OAuthConfig Timeout int UseStreaming bool // Whether to use streaming HTTP for better performance SkipVerify bool // Whether to skip TLS certificate verification for HTTPS connections Headers map[string]string // Custom headers to include in HTTP requests // No LLM configuration needed - MCP doesn't need to know about LLM models } // AuthConfig represents authentication options for HTTP MCP servers type AuthConfig struct { Type string // "none", "basic", "bearer", "api-key" Username string // For basic auth Password string // For basic auth Token string // For bearer auth ApiKey string // For API key auth HeaderName string // Custom header name for API key } // OAuthConfig represents OAuth configuration for HTTP MCP servers type OAuthConfig struct { ClientID string ClientSecret string TokenURL string AuthURL string Scopes []string RedirectURL string } // NewMCPClient creates a new MCP client with the appropriate implementation based on the config func NewMCPClient(config ClientConfig) (MCPClient, error) { // Validate common configuration if config.Name == "" { return nil, fmt.Errorf("client name is required") } // Choose the appropriate client implementation if config.URL != "" { // Use HTTP client return NewHTTPClient(config), nil } // Default to stdio client if config.Command == "" { return nil, fmt.Errorf("either URL or Command must be specified") } return NewStdioClient(config), nil } ================================================ FILE: pkg/mcp/manager.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "context" "fmt" "sync" "time" "k8s.io/klog/v2" ) // ============================================================================= // Status Types // ============================================================================= // ServerConnectionInfo holds connection status for a single MCP server type ServerConnectionInfo struct { Name string Command string IsLegacy bool IsConnected bool AvailableTools []Tool } // MCPStatus represents the overall status of MCP servers and tools type MCPStatus struct { ServerInfoList []ServerConnectionInfo TotalServers int ConnectedCount int FailedCount int TotalTools int ClientEnabled bool } // ============================================================================= // Manager Core // ============================================================================= // Manager handles MCP client connections and tool discovery type Manager struct { config *Config clients map[string]*Client mu sync.RWMutex } // NewManager creates a new MCP manager with the given configuration func NewManager(config *Config) *Manager { return &Manager{ config: config, clients: make(map[string]*Client), } } // InitializeManager creates and initializes the MCP manager // with configuration loaded from default paths func InitializeManager() (*Manager, error) { klog.V(1).Info("Initializing MCP client functionality") config, err := LoadConfig("") if err != nil { klog.V(2).Info("Failed to load MCP config", "error", err) return nil, err } return NewManager(config), nil } // ============================================================================= // Connection Management // ============================================================================= // ConnectAll connects to all configured MCP servers func (m *Manager) ConnectAll(ctx context.Context) error { m.mu.Lock() defer m.mu.Unlock() var errs []error for _, serverCfg := range m.config.Servers { if _, exists := m.clients[serverCfg.Name]; exists { klog.V(2).Info("MCP client already connected", "name", serverCfg.Name) continue } // Convert environment map to slice var envSlice []string for k, v := range serverCfg.Env { envSlice = append(envSlice, fmt.Sprintf("%s=%s", k, v)) } // Create client config with environment map config := ClientConfig{ Name: serverCfg.Name, Command: serverCfg.Command, Args: serverCfg.Args, Auth: serverCfg.Auth, OAuthConfig: serverCfg.OAuthConfig, Env: envSlice, URL: serverCfg.URL, Timeout: serverCfg.Timeout, UseStreaming: serverCfg.UseStreaming, SkipVerify: serverCfg.SkipVerify, } client := NewClient(config) if err := client.Connect(ctx); err != nil { err := fmt.Errorf(ErrServerConnectionFmt, serverCfg.Name, err) errs = append(errs, err) klog.Error(err) continue } m.clients[serverCfg.Name] = client klog.V(2).Info("Connected to MCP server", "name", serverCfg.Name) } if len(errs) > 0 { return fmt.Errorf("failed to connect to some MCP servers: %v", errs) } return nil } // Close closes all MCP client connections func (m *Manager) Close() error { m.mu.Lock() defer m.mu.Unlock() var errs []error for name, client := range m.clients { if err := client.Close(); err != nil { errs = append(errs, fmt.Errorf(ErrServerCloseFmt, name, err)) } delete(m.clients, name) } if len(errs) > 0 { return fmt.Errorf("errors while closing MCP clients: %v", errs) } return nil } // GetClient returns a connected MCP client by name func (m *Manager) GetClient(name string) (*Client, bool) { m.mu.RLock() defer m.mu.RUnlock() client, exists := m.clients[name] return client, exists } // ListClients returns a list of all connected MCP clients func (m *Manager) ListClients() []*Client { m.mu.RLock() defer m.mu.RUnlock() var clients []*Client for _, client := range m.clients { clients = append(clients, client) } return clients } // ============================================================================= // Server and Tool Discovery // ============================================================================= // DiscoverAndConnectServers connects to all configured servers // with a timeout and stabilization delay func (m *Manager) DiscoverAndConnectServers(ctx context.Context) error { klog.V(1).Info("Connecting to MCP servers") connectCtx, connectCancel := context.WithTimeout(ctx, DefaultConnectionTimeout) defer connectCancel() if err := m.ConnectAll(connectCtx); err != nil { klog.V(2).Info("Failed to connect to some MCP servers during auto-discovery", "error", err) // Continue with partial connections } // Allow connections to stabilize before tool discovery klog.V(3).Info("Waiting for server connections to stabilize", "delay", DefaultStabilizationDelay) time.Sleep(DefaultStabilizationDelay) return nil } // ListAvailableTools returns tools from all connected servers // For retries and more robust handling, use RefreshToolDiscovery func (m *Manager) ListAvailableTools(ctx context.Context) (map[string][]Tool, error) { m.mu.RLock() defer m.mu.RUnlock() tools := make(map[string][]Tool) for name, client := range m.clients { toolList, err := client.ListTools(ctx) if err != nil { return nil, fmt.Errorf("listing tools from MCP server %q: %w", name, err) } var serverTools []Tool for _, tool := range toolList { serverTools = append(serverTools, tool.WithServer(name)) } tools[name] = serverTools } return tools, nil } // RefreshToolDiscovery discovers tools from all servers with retries func (m *Manager) RefreshToolDiscovery(ctx context.Context) (map[string][]Tool, error) { klog.V(1).Info("Starting tool discovery from MCP servers with retries") var serverTools map[string][]Tool retryConfig := DefaultRetryConfig("tool discovery from MCP servers") err := RetryOperation(ctx, retryConfig, func() error { var err error serverTools, err = m.ListAvailableTools(ctx) return err }) if err != nil { klog.Warningf("Failed to discover tools after retries: %v", err) return nil, err } // Log discovery results toolCount := 0 for serverName, tools := range serverTools { klog.V(1).Info("Discovered tools from MCP server", "server", serverName, "toolCount", len(tools)) toolCount += len(tools) } if toolCount > 0 { klog.InfoS("Successfully discovered MCP tools", "totalTools", toolCount) } else { klog.V(1).Info("No MCP tools were discovered from connected servers") } return serverTools, nil } // RegisterTools discovers and registers tools from all MCP servers using the provided callback // The callback function is responsible for creating and registering tool wrappers func (m *Manager) RegisterTools(ctx context.Context, registerCallback func(serverName string, tool Tool) error) error { // Discover tools from connected servers serverTools, err := m.RefreshToolDiscovery(ctx) if err != nil { return err } toolCount := 0 for serverName, tools := range serverTools { for _, toolInfo := range tools { // Use the callback to register each tool if err := registerCallback(serverName, toolInfo); err != nil { klog.Warningf("Failed to register tool %s from server %s: %v", toolInfo.Name, serverName, err) continue } toolCount++ } } if toolCount > 0 { klog.InfoS("Registered MCP tools", "totalTools", toolCount) } return nil } // ============================================================================= // Status Reporting // ============================================================================= // GetStatus returns status of all MCP servers and their tools func (m *Manager) GetStatus(ctx context.Context, mcpClientEnabled bool) (*MCPStatus, error) { status := &MCPStatus{ ClientEnabled: mcpClientEnabled, } mcpConfigPath, err := DefaultConfigPath() if err != nil { klog.V(2).Infof("Failed to get MCP config path: %v", err) return status, nil // Return empty status } mcpConfig, err := LoadConfig(mcpConfigPath) if err != nil { return status, nil // Return empty status } status.TotalServers = len(mcpConfig.Servers) if status.TotalServers == 0 { return status, nil } var serverTools map[string][]Tool var connectedClients []*Client if mcpClientEnabled && m != nil { connectedClients = m.ListClients() status.ConnectedCount = len(connectedClients) status.FailedCount = status.TotalServers - status.ConnectedCount toolsCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() serverTools, err = m.ListAvailableTools(toolsCtx) if err != nil { klog.V(2).InfoS("Failed to get tools from MCP manager", "error", err) serverTools = make(map[string][]Tool) } for _, toolList := range serverTools { status.TotalTools += len(toolList) } } else { serverTools = make(map[string][]Tool) } connectedServerNames := make(map[string]bool) if mcpClientEnabled { for _, client := range connectedClients { connectedServerNames[client.Name] = true } } // Process all servers for _, server := range mcpConfig.Servers { serverInfo := ServerConnectionInfo{ Name: server.Name, Command: server.Command, IsLegacy: false, IsConnected: connectedServerNames[server.Name], } if tools, exists := serverTools[server.Name]; exists { serverInfo.AvailableTools = tools } status.ServerInfoList = append(status.ServerInfoList, serverInfo) } return status, nil } // LogConfig logs the MCP configuration summary // If mcpConfigPath is empty, uses the Manager's existing config func (m *Manager) LogConfig(mcpConfigPath string) error { var mcpConfig *Config var err error if mcpConfigPath == "" && m.config != nil { mcpConfig = m.config } else { mcpConfig, err = LoadConfig(mcpConfigPath) if err != nil { return fmt.Errorf("failed to load MCP config from %s: %w", mcpConfigPath, err) } } serverCount := len(mcpConfig.Servers) totalServers := serverCount if totalServers > 0 { serverWord := "server" if totalServers > 1 { serverWord = "servers" } if mcpConfigPath != "" { klog.V(2).Infof("Loaded %d MCP %s from %s", totalServers, serverWord, mcpConfigPath) } else { klog.V(2).Infof("Found %d MCP %s in configuration", totalServers, serverWord) } for _, server := range mcpConfig.Servers { klog.V(2).Infof(" - %s: %s", server.Name, server.Command) } } return nil } // ============================================================================= // Integration Methods // ============================================================================= // RegisterWithToolSystem connects to MCP servers and registers discovered tools with an external tool system // using the provided callback function. This simplifies integration with kubectl-ai's tool system. func (m *Manager) RegisterWithToolSystem(ctx context.Context, registerCallback func(serverName string, tool Tool) error) error { klog.V(1).Info("Initializing MCP client functionality and registering tools") // Connect to all configured servers if err := m.DiscoverAndConnectServers(ctx); err != nil { return fmt.Errorf("MCP server connection failed: %w", err) } // Register all discovered tools using the callback if err := m.RegisterTools(ctx, registerCallback); err != nil { return fmt.Errorf("MCP tool registration failed: %w", err) } return nil } ================================================ FILE: pkg/mcp/stdio_client.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "context" "fmt" mcpclient "github.com/mark3labs/mcp-go/client" mcp "github.com/mark3labs/mcp-go/mcp" "k8s.io/klog/v2" ) // =================================================================== // Stdio Client Implementation // =================================================================== // stdioClient is an MCP client that communicates via standard I/O type stdioClient struct { name string command string args []string env []string client *mcpclient.Client } // NewStdioClient creates a new stdio-based MCP client func NewStdioClient(config ClientConfig) MCPClient { return &stdioClient{ name: config.Name, command: config.Command, args: config.Args, env: config.Env, } } // getUnderlyingClient returns the underlying MCP client. func (c *stdioClient) getUnderlyingClient() *mcpclient.Client { return c.client } // ensureConnected makes sure the client is connected. func (c *stdioClient) ensureConnected() error { return ensureClientConnected(c.client) } // Name returns the name of this client. func (c *stdioClient) Name() string { return c.name } // Connect establishes a connection to the stdio MCP server. func (c *stdioClient) Connect(ctx context.Context) error { klog.V(2).InfoS("Connecting to stdio MCP server", "name", c.name, "command", c.command) if c.client != nil { return nil // Already connected } // Expand the command path and prepare the environment expandedCmd, err := expandPath(c.command) if err != nil { return fmt.Errorf("expanding command path: %w", err) } // Create the stdio MCP client client, err := mcpclient.NewStdioMCPClient(expandedCmd, c.env, c.args...) if err != nil { return fmt.Errorf("creating stdio MCP client: %w", err) } c.client = client // Initialize the connection if err := c.initializeConnection(ctx); err != nil { c.cleanup() return fmt.Errorf("initializing connection: %w", err) } // Verify the connection if err := c.verifyConnection(ctx); err != nil { c.cleanup() return fmt.Errorf("verifying connection: %w", err) } klog.V(2).InfoS("Successfully connected to stdio MCP server", "name", c.name) return nil } // initializeConnection initializes the MCP connection with proper handshake func (c *stdioClient) initializeConnection(ctx context.Context) error { return initializeClientConnection(ctx, c.client) } // verifyConnection verifies the connection works by testing tool listing func (c *stdioClient) verifyConnection(ctx context.Context) error { return verifyClientConnection(ctx, c.client) } // cleanup closes the client connection and resets the client state func (c *stdioClient) cleanup() { cleanupClient(&c.client) } // Close closes the connection to the MCP server func (c *stdioClient) Close() error { if c.client == nil { return nil // Already closed } klog.V(2).InfoS("Closing connection to stdio MCP server", "name", c.name) err := c.client.Close() c.client = nil if err != nil { return fmt.Errorf("closing MCP client: %w", err) } return nil } // ListTools lists all available tools from the MCP server func (c *stdioClient) ListTools(ctx context.Context) ([]Tool, error) { tools, err := listClientTools(ctx, c.client, c.name) if err != nil { return nil, err } klog.V(2).InfoS("Listed tools from stdio MCP server", "count", len(tools), "server", c.name) return tools, nil } // CallTool calls a tool on the MCP server and returns the result as a string func (c *stdioClient) CallTool(ctx context.Context, toolName string, arguments map[string]interface{}) (string, error) { klog.V(2).InfoS("Calling MCP tool via stdio", "server", c.name, "tool", toolName) if err := c.ensureConnected(); err != nil { return "", err } // Create v0.31.0 compatible request request := mcp.CallToolRequest{ Params: mcp.CallToolParams{ Name: toolName, Arguments: arguments, }, } // Call the tool on the MCP server result, err := c.client.CallTool(ctx, request) if err != nil { return "", fmt.Errorf("error calling tool %s: %w", toolName, err) } return processToolResponse(result) } ================================================ FILE: pkg/mcp/utils.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package mcp import ( "context" "fmt" "math" "os" "os/exec" "path/filepath" "strconv" "strings" "time" "k8s.io/klog/v2" ) // RetryConfig defines retry behavior for MCP operations type RetryConfig struct { MaxRetries int BaseDelay time.Duration MaxDelay time.Duration Multiplier float64 Description string } // DefaultRetryConfig returns a sensible default retry configuration func DefaultRetryConfig(description string) RetryConfig { return RetryConfig{ MaxRetries: 3, BaseDelay: 1 * time.Second, MaxDelay: 10 * time.Second, Multiplier: 2.0, Description: description, } } // RetryOperation executes an operation with exponential backoff retry func RetryOperation(ctx context.Context, config RetryConfig, operation func() error) error { var lastErr error for attempt := 1; attempt <= config.MaxRetries; attempt++ { klog.V(3).InfoS("Attempting operation", "operation", config.Description, "attempt", attempt, "maxRetries", config.MaxRetries) if err := operation(); err == nil { if attempt > 1 { klog.V(2).InfoS("Operation succeeded after retry", "operation", config.Description, "attempt", attempt) } return nil } else { lastErr = err if attempt < config.MaxRetries { delay := calculateBackoffDelay(attempt, config) klog.V(3).InfoS("Operation failed, retrying", "operation", config.Description, "attempt", attempt, "error", err, "nextRetryIn", delay) select { case <-ctx.Done(): return fmt.Errorf("operation cancelled: %w", ctx.Err()) case <-time.After(delay): // Continue to next attempt } } } } return fmt.Errorf("operation failed after %d attempts: %w", config.MaxRetries, lastErr) } // calculateBackoffDelay calculates exponential backoff delay with jitter func calculateBackoffDelay(attempt int, config RetryConfig) time.Duration { delay := float64(config.BaseDelay) * math.Pow(config.Multiplier, float64(attempt-1)) if time.Duration(delay) > config.MaxDelay { return config.MaxDelay } return time.Duration(delay) } // SanitizeServerName ensures server names are valid identifiers func SanitizeServerName(name string) string { // Simple sanitization - replace invalid characters result := "" for _, char := range name { if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') || char == '-' || char == '_' { result += string(char) } else { result += "_" } } if result == "" { result = "unnamed" } return result } // GroupToolsByServer groups tools by their server name for easier display func GroupToolsByServer(tools map[string][]Tool) map[string]int { summary := make(map[string]int) for serverName, serverTools := range tools { summary[serverName] = len(serverTools) } return summary } // mergeEnvironmentVariables merges process environment with custom environment variables func mergeEnvironmentVariables(processEnv, customEnv []string) []string { envMap := make(map[string]string) // Parse process environment for _, e := range processEnv { if parts := strings.SplitN(e, "=", 2); len(parts) == 2 { envMap[parts[0]] = parts[1] } } // Override with custom environment variables for _, env := range customEnv { if parts := strings.SplitN(env, "=", 2); len(parts) == 2 { envMap[parts[0]] = parts[1] } } // Convert back to slice finalEnv := make([]string, 0, len(envMap)) for k, v := range envMap { finalEnv = append(finalEnv, fmt.Sprintf("%s=%s", k, v)) } return finalEnv } // expandPath expands the command path, handling ~ and environment variables // If the path is just a binary name (no path separators), it looks in $PATH func expandPath(path string) (string, error) { if path == "" { return "", fmt.Errorf("path cannot be empty") } // Expand environment variables first expanded := os.ExpandEnv(path) // If the command contains no path separators, look it up in $PATH first if !strings.Contains(expanded, string(filepath.Separator)) && !strings.HasPrefix(expanded, "~") { klog.V(2).InfoS("Attempting PATH lookup for command", "command", expanded) // Try to find the command in $PATH if pathResolved, err := exec.LookPath(expanded); err == nil { klog.V(2).InfoS("Found command in PATH", "command", expanded, "resolved", pathResolved) return pathResolved, nil } else { klog.V(2).InfoS("Command not found in PATH", "command", expanded, "error", err) } // If not found in PATH, continue with the original logic below klog.V(2).InfoS("Command not found in PATH, trying relative to current directory", "command", expanded) } else { klog.V(2).InfoS("Skipping PATH lookup", "command", expanded, "hasPathSeparator", strings.Contains(expanded, string(filepath.Separator)), "hasTilde", strings.HasPrefix(expanded, "~")) } // Handle ~ for home directory if strings.HasPrefix(expanded, "~") { home, err := os.UserHomeDir() if err != nil { return "", fmt.Errorf("getting home directory: %w", err) } expanded = filepath.Join(home, expanded[1:]) } // Clean the path to remove any . or .. elements expanded = filepath.Clean(expanded) // Make the path absolute if it's not already if !filepath.IsAbs(expanded) { cwd, err := os.Getwd() if err != nil { return "", fmt.Errorf("getting current working directory: %w", err) } expanded = filepath.Clean(filepath.Join(cwd, expanded)) } // Verify the file exists and is executable info, err := os.Stat(expanded) if err != nil { return "", fmt.Errorf(ErrPathCheckFmt, expanded, err) } // Check if it's a regular file and executable if !info.Mode().IsRegular() { return "", fmt.Errorf("path %q is not a regular file", expanded) } // Check if the file is executable by the current user if info.Mode().Perm()&0111 == 0 { return "", fmt.Errorf("file %q is not executable", expanded) } return expanded, nil } // ============================================================================= // Helper Functions to Reduce Redundancy // ============================================================================= // withTimeout creates a context with the specified timeout and returns the context and cancel function func withTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) { return context.WithTimeout(ctx, timeout) } // ensureConnected checks if the client is connected and returns an error if not func (c *Client) ensureConnected() error { if c.client == nil { return fmt.Errorf("not connected to MCP server") } return nil } // ============================================================================= // MCP Tool Helper Functions // ============================================================================= // FunctionDefinition is an interface representing generic function schema definitions // This allows the MCP package to create schemas without directly depending on gollm type FunctionDefinition interface { // Schema returns a representation of the function schema Schema() any } // SchemaProperty is an interface representing generic schema properties type SchemaProperty interface { // Property returns a representation of the schema property Property() any } // SchemaBuilder is a function that builds a function definition from a tool type SchemaBuilder func(tool *Tool) (FunctionDefinition, error) // ConvertArgs handles all argument conversions for MCP tools. // It transforms keys from snake_case to camelCase and converts values to appropriate types. func ConvertArgs(args map[string]any) map[string]any { if len(args) == 0 { return args } result := make(map[string]any, len(args)) for key, value := range args { // Convert key from snake_case to camelCase camelKey := SnakeToCamel(key) // Convert value based on key name patterns result[camelKey] = ConvertValue(camelKey, value) } return result } // SnakeToCamel converts a snake_case string to camelCase. func SnakeToCamel(s string) string { if !strings.Contains(s, "_") { return s } parts := strings.Split(s, "_") result := parts[0] for _, part := range parts[1:] { if len(part) > 0 { result += strings.ToUpper(part[:1]) + part[1:] } } return result } // ConvertValue infers and converts a value to an appropriate type based on the parameter name. func ConvertValue(paramName string, value any) any { // Already primitive types that don't need conversion switch value.(type) { case bool, int, int32, int64, float32, float64: return value } name := strings.ToLower(paramName) // Number parameter detection if IsNumberParam(name) { if str, ok := value.(string); ok { // Try integer conversion first if num, err := strconv.Atoi(str); err == nil { return num } // Then try float conversion if num, err := strconv.ParseFloat(str, 64); err == nil { return num } } else if f, ok := value.(float64); ok && f == float64(int(f)) { // Convert whole number floats to int return int(f) } } // Boolean parameter detection if IsBoolParam(name) { if str, ok := value.(string); ok { if b, err := strconv.ParseBool(str); err == nil { return b } } else if n, ok := value.(int); ok { return n != 0 } } return value } // IsNumberParam checks if a parameter name suggests a numeric value. func IsNumberParam(name string) bool { numberPatterns := []string{"number", "count", "total", "max", "min", "limit"} for _, pattern := range numberPatterns { if strings.Contains(name, pattern) { return true } } return false } // IsBoolParam checks if a parameter name suggests a boolean value. func IsBoolParam(name string) bool { // Prefix checks boolPrefixes := []string{"is", "has", "needs", "enable"} for _, prefix := range boolPrefixes { if strings.HasPrefix(name, prefix) { return true } } // Contains checks boolPatterns := []string{"required", "enabled"} for _, pattern := range boolPatterns { if strings.Contains(name, pattern) { return true } } return false } ================================================ FILE: pkg/sandbox/executor.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package sandbox import ( "context" "fmt" ) // Executor defines the interface for executing commands. type Executor interface { // Execute runs a command and returns the result. Execute(ctx context.Context, command string, env []string, workDir string) (*ExecResult, error) // Close cleans up any resources associated with the executor. Close(ctx context.Context) error } // ExecResult represents the result of a command execution. type ExecResult struct { Command string `json:"command,omitempty"` Error string `json:"error,omitempty"` Stdout string `json:"stdout,omitempty"` Stderr string `json:"stderr,omitempty"` ExitCode int `json:"exit_code,omitempty"` StreamType string `json:"stream_type,omitempty"` } func (e *ExecResult) String() string { return fmt.Sprintf("Command: %q\nError: %q\nStdout: %q\nStderr: %q\nExitCode: %d\nStreamType: %q}", e.Command, e.Error, e.Stdout, e.Stderr, e.ExitCode, e.StreamType) } ================================================ FILE: pkg/sandbox/kubernetes.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Package sandbox provides Kubernetes-based sandboxed command execution with an exec.Command-like interface // A sandbox represents an isolated execution environment. Currently implemented using Kubernetes pods, // but can be extended to support other backends like Docker containers in the future. package sandbox import ( "bytes" "context" "fmt" "io" "os" "strings" "time" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/util/wait" "k8s.io/client-go/kubernetes" "k8s.io/client-go/kubernetes/scheme" "k8s.io/client-go/rest" "k8s.io/client-go/tools/clientcmd" "k8s.io/client-go/tools/remotecommand" ) // KubernetesSandbox represents a Kubernetes-based sandboxed execution environment type KubernetesSandbox struct { name string namespace string image string kubeconfig string clientset *kubernetes.Clientset config *rest.Config } // Execute executes the command in the sandbox. func (s *KubernetesSandbox) Execute(ctx context.Context, command string, env []string, workDir string) (*ExecResult, error) { fullCommand := command // Ensure kubectl is in the PATH fullCommand = fmt.Sprintf("export PATH=/opt/bitnami/kubectl/bin:$PATH; %s", fullCommand) if workDir != "" { fullCommand = fmt.Sprintf("mkdir -p %q && cd %q && %s", workDir, workDir, fullCommand) } for _, envVar := range env { fullCommand = fmt.Sprintf("export %s; %s", envVar, fullCommand) } cmd := s.CommandContext(ctx, fullCommand) output, err := cmd.CombinedOutput() result := &ExecResult{ Command: command, Stdout: string(output), } if err != nil { result.Error = err.Error() result.ExitCode = 1 } return result, nil } // Close cleans up the sandbox resources. func (s *KubernetesSandbox) Close(ctx context.Context) error { return s.Delete(ctx) } // Cmd represents a command to be executed in a sandbox // It follows the same interface pattern as exec.Cmd type Cmd struct { sandbox *KubernetesSandbox command []string ctx context.Context // Standard streams (similar to exec.Cmd) Stdin io.Reader Stdout io.Writer Stderr io.Writer } // Option represents a configuration option for KubernetesSandbox type Option func(*KubernetesSandbox) error // NewKubernetesSandbox creates a new KubernetesSandbox instance with the given name and options func NewKubernetesSandbox(name string, opts ...Option) (*KubernetesSandbox, error) { s := &KubernetesSandbox{ name: name, namespace: "computer", // default namespace } // Apply options for _, opt := range opts { if err := opt(s); err != nil { return nil, err } } // Initialize Kubernetes client config, err := clientcmd.BuildConfigFromFlags("", s.kubeconfig) if err != nil { return nil, fmt.Errorf("error building kubeconfig: %v", err) } clientset, err := kubernetes.NewForConfig(config) if err != nil { return nil, fmt.Errorf("error creating Kubernetes client: %v", err) } s.config = config s.clientset = clientset return s, nil } // WithKubeconfig sets the kubeconfig file path func WithKubeconfig(kubeconfig string) Option { return func(s *KubernetesSandbox) error { s.kubeconfig = kubeconfig return nil } } // WithName sets the sandbox name (deprecated - use constructor parameter instead) func WithName(name string) Option { return func(s *KubernetesSandbox) error { s.name = name return nil } } // WithNamespace sets the namespace func WithNamespace(namespace string) Option { return func(s *KubernetesSandbox) error { s.namespace = namespace return nil } } // WithImage sets the container image func WithImage(image string) Option { return func(s *KubernetesSandbox) error { s.image = image return nil } } // Command creates a new Cmd to execute the given command in the sandbox // This follows the same interface as exec.Command func (s *KubernetesSandbox) Command(name string, arg ...string) *Cmd { cmd := &Cmd{ sandbox: s, command: append([]string{name}, arg...), ctx: context.Background(), } return cmd } // CommandContext creates a new Cmd with a context func (s *KubernetesSandbox) CommandContext(ctx context.Context, name string, arg ...string) *Cmd { cmd := &Cmd{ sandbox: s, command: append([]string{name}, arg...), ctx: ctx, } return cmd } // Delete removes the sandbox pod and its associated resources, waiting for them to be fully terminated. // It does not return an error if the resources are already deleted. func (s *KubernetesSandbox) Delete(ctx context.Context) error { var errs []string // 1. Initiate deletion of the Pod with a zero grace period for faster removal. deleteOptions := metav1.DeleteOptions{ GracePeriodSeconds: new(int64), // 0 seconds } err := s.clientset.CoreV1().Pods(s.namespace).Delete(ctx, s.name, deleteOptions) if err != nil && !errors.IsNotFound(err) { errs = append(errs, fmt.Sprintf("failed to initiate pod deletion: %v", err)) } // 2. Initiate deletion of the ConfigMap. configMapName := s.name + "-kubeconfig" if err := s.deleteKubeconfigMap(ctx, configMapName); err != nil { errs = append(errs, fmt.Sprintf("failed to initiate configmap deletion: %v", err)) } // 3. Wait for the Pod to be fully terminated. pollErr := wait.PollUntilContextTimeout(ctx, 2*time.Second, 1*time.Minute, true, func(ctx context.Context) (bool, error) { _, getErr := s.clientset.CoreV1().Pods(s.namespace).Get(ctx, s.name, metav1.GetOptions{}) if errors.IsNotFound(getErr) { return true, nil // Pod is gone. } if getErr != nil { return false, getErr // Polling failed with an unexpected error. } return false, nil // Pod still exists, continue polling. }) if pollErr != nil { errs = append(errs, fmt.Sprintf("error waiting for pod deletion: %v", pollErr)) } if len(errs) > 0 { return fmt.Errorf("errors during sandbox deletion: %s", strings.Join(errs, "; ")) } return nil } // Run executes the command and waits for it to complete func (c *Cmd) Run() error { return c.execute(nil, nil) } // Output runs the command and returns its standard output func (c *Cmd) Output() ([]byte, error) { var stdout bytes.Buffer err := c.execute(&stdout, nil) return stdout.Bytes(), err } // CombinedOutput runs the command and returns its combined standard output and standard error func (c *Cmd) CombinedOutput() ([]byte, error) { var output bytes.Buffer err := c.execute(&output, &output) return output.Bytes(), err } // execute is the internal method that handles the actual pod execution func (c *Cmd) execute(stdout, stderr io.Writer) error { sandbox := c.sandbox // Validate required fields if sandbox.name == "" || sandbox.image == "" { return fmt.Errorf("sandbox name and image must be specified") } // Check if pod exists and validate its image if it does. existingPod, err := c.getPod() if err != nil { return fmt.Errorf("error checking for existing sandbox: %w", err) } if existingPod != nil { // Sandbox exists. Verify the container image matches. var existingImage string for _, container := range existingPod.Spec.Containers { if container.Name == "main" { existingImage = container.Image break } } if existingImage != "" && existingImage != sandbox.image { return fmt.Errorf( "existing sandbox '%s' uses image '%s', but new execution requested image '%s'. Please delete the sandbox first", sandbox.name, existingImage, sandbox.image, ) } } else { // Pod doesn't exist, create it. if err := c.createPod(); err != nil { return fmt.Errorf("error creating pod: %v", err) } } // Wait for pod to be ready if err := c.waitForPodReady(); err != nil { return fmt.Errorf("error waiting for pod to be ready: %v", err) } // Execute command in pod return c.executeInPod(stdout, stderr) } // getPod fetches the sandbox pod if it exists. Returns (nil, nil) if not found. func (c *Cmd) getPod() (*corev1.Pod, error) { sandbox := c.sandbox pod, err := sandbox.clientset.CoreV1().Pods(sandbox.namespace).Get(c.ctx, sandbox.name, metav1.GetOptions{}) if err != nil { if errors.IsNotFound(err) { return nil, nil // Not an error, just means we need to create it. } return nil, err } return pod, nil } // createPod creates a new pod for the sandbox, including its kubeconfig configmap func (c *Cmd) createPod() error { sandbox := c.sandbox configMapName := sandbox.name + "-kubeconfig" // Create a dedicated kubeconfig for the pod to use. // This ensures kubectl defaults to the "default" namespace. if err := c.createKubeconfigMap(configMapName); err != nil { // If the configmap already exists, we can proceed. if !errors.IsAlreadyExists(err) { return fmt.Errorf("failed to create in-pod kubeconfig: %w", err) } } pod := &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Name: sandbox.name, Namespace: sandbox.namespace, }, Spec: corev1.PodSpec{ ServiceAccountName: "normal-user", Containers: []corev1.Container{ { Name: "main", Image: sandbox.image, Command: []string{"sleep"}, Args: []string{"infinity"}, // Sleep forever to keep the container running Env: []corev1.EnvVar{ { Name: "KUBECONFIG", Value: "/etc/kube/config", }, { Name: "PATH", Value: "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/bitnami/kubectl/bin", }, }, VolumeMounts: []corev1.VolumeMount{ { Name: "kubeconfig-volume", MountPath: "/etc/kube", ReadOnly: true, }, }, }, }, Volumes: []corev1.Volume{ { Name: "kubeconfig-volume", VolumeSource: corev1.VolumeSource{ ConfigMap: &corev1.ConfigMapVolumeSource{ LocalObjectReference: corev1.LocalObjectReference{ Name: configMapName, }, Items: []corev1.KeyToPath{ { Key: "config", Path: "config", }, }, }, }, }, }, RestartPolicy: corev1.RestartPolicyNever, }, } _, podCreateErr := sandbox.clientset.CoreV1().Pods(sandbox.namespace).Create(c.ctx, pod, metav1.CreateOptions{}) if podCreateErr != nil { // If pod creation fails, attempt to clean up the configmap we just created. if cleanupErr := sandbox.deleteKubeconfigMap(c.ctx, configMapName); cleanupErr != nil { return fmt.Errorf("pod creation failed: %v; ALSO, configmap cleanup failed: %v", podCreateErr, cleanupErr) } return fmt.Errorf("pod creation failed: %w", podCreateErr) } return nil } // createKubeconfigMap generates a kubeconfig file that uses the pod's service account token // and sets the default namespace to "default". This is stored in a ConfigMap. func (c *Cmd) createKubeconfigMap(name string) error { sandbox := c.sandbox // Use a static string template for the kubeconfig to ensure correctness. kubeconfigYAML := fmt.Sprintf(`apiVersion: v1 clusters: - cluster: server: https://kubernetes.default.svc certificate-authority: /var/run/secrets/kubernetes.io/serviceaccount/ca.crt name: default contexts: - context: cluster: default namespace: %s user: default name: default current-context: default users: - name: default user: tokenFile: /var/run/secrets/kubernetes.io/serviceaccount/token`, sandbox.namespace) // Create the ConfigMap object. configMap := &corev1.ConfigMap{ ObjectMeta: metav1.ObjectMeta{ Name: name, Namespace: sandbox.namespace, }, Data: map[string]string{ "config": kubeconfigYAML, }, } _, err := sandbox.clientset.CoreV1().ConfigMaps(sandbox.namespace).Create(c.ctx, configMap, metav1.CreateOptions{}) return err } // deleteKubeconfigMap cleans up the ConfigMap created for the pod. func (s *KubernetesSandbox) deleteKubeconfigMap(ctx context.Context, name string) error { err := s.clientset.CoreV1().ConfigMaps(s.namespace).Delete(ctx, name, metav1.DeleteOptions{}) if err != nil && !errors.IsNotFound(err) { return fmt.Errorf("failed to delete kubeconfig configmap: %w", err) } return nil } // waitForPodReady waits for the pod to be ready func (c *Cmd) waitForPodReady() error { sandbox := c.sandbox return wait.PollUntilContextTimeout(c.ctx, 2*time.Second, 5*time.Minute, true, func(ctx context.Context) (bool, error) { pod, err := sandbox.clientset.CoreV1().Pods(sandbox.namespace).Get(ctx, sandbox.name, metav1.GetOptions{}) if err != nil { return false, err } // Check if pod is ready for _, condition := range pod.Status.Conditions { if condition.Type == corev1.PodReady && condition.Status == corev1.ConditionTrue { return true, nil } } // Check if pod failed if pod.Status.Phase == corev1.PodFailed { return false, fmt.Errorf("pod %s failed", sandbox.name) } return false, nil }) } // executeInPod executes the command in the pod func (c *Cmd) executeInPod(stdout, stderr io.Writer) error { sandbox := c.sandbox // Use provided writers or default to the Cmd's streams if stdout == nil { stdout = c.Stdout if stdout == nil { stdout = os.Stdout } } if stderr == nil { stderr = c.Stderr if stderr == nil { stderr = os.Stderr } } req := sandbox.clientset.CoreV1().RESTClient().Post(). Resource("pods"). Name(sandbox.name). Namespace(sandbox.namespace). SubResource("exec") commandStr := strings.Join(c.command, " ") req.VersionedParams(&corev1.PodExecOptions{ Container: "main", Command: []string{"/bin/sh", "-c", commandStr}, Stdin: c.Stdin != nil, Stdout: true, Stderr: true, TTY: false, }, scheme.ParameterCodec) exec, err := remotecommand.NewSPDYExecutor(sandbox.config, "POST", req.URL()) if err != nil { return fmt.Errorf("error creating executor: %v", err) } err = exec.StreamWithContext(c.ctx, remotecommand.StreamOptions{ Stdin: c.Stdin, Stdout: stdout, Stderr: stderr, }) if err != nil { return fmt.Errorf("error executing command: %v", err) } return nil } ================================================ FILE: pkg/sandbox/local.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package sandbox import ( "bytes" "context" "os" "os/exec" "runtime" "k8s.io/klog/v2" ) const ( defaultBashBin = "/bin/bash" ) // Local executes commands locally. type Local struct{} // NewLocalExecutor creates a new LocalExecutor. func NewLocalExecutor() *Local { return &Local{} } // Execute executes the command locally. func (e *Local) Execute(ctx context.Context, command string, env []string, workDir string) (*ExecResult, error) { // Use the provided context directly cmdCtx := ctx var cmd *exec.Cmd if runtime.GOOS == "windows" { cmd = exec.CommandContext(cmdCtx, os.Getenv("COMSPEC"), "/c", command) } else { cmd = exec.CommandContext(cmdCtx, lookupBashBin(), "-c", command) } cmd.Dir = workDir cmd.Env = env var stdoutBuf, stderrBuf bytes.Buffer cmd.Stdout = &stdoutBuf cmd.Stderr = &stderrBuf err := cmd.Run() result := &ExecResult{ Command: command, Stdout: stdoutBuf.String(), Stderr: stderrBuf.String(), } if err != nil { // If it wasn't a timeout (or not a streaming command), it's a real error if exitError, ok := err.(*exec.ExitError); ok { result.ExitCode = exitError.ExitCode() result.Error = exitError.Error() // Stderr is already captured in result.Stderr } else { return nil, err } } return result, nil } // Close is a no-op for Local executor. func (e *Local) Close(ctx context.Context) error { return nil } // Find the bash executable path using exec.LookPath. func lookupBashBin() string { actualBashPath, err := exec.LookPath("bash") if err != nil { klog.Warningf("'bash' not found in PATH, defaulting to %s: %v", defaultBashBin, err) return defaultBashBin } return actualBashPath } ================================================ FILE: pkg/sandbox/seatbelt_executor.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //go:build darwin package sandbox import ( "bytes" "context" "fmt" "os/exec" ) // Seatbelt executes commands in a seatbelt sandbox. type Seatbelt struct { local *Local } // NewSeatbeltExecutor creates a new SeatbeltExecutor. func NewSeatbeltExecutor() *Seatbelt { return &Seatbelt{ local: NewLocalExecutor(), } } // Execute executes the command in the seatbelt sandbox. func (e *Seatbelt) Execute(ctx context.Context, command string, env []string, workDir string) (*ExecResult, error) { // Use the provided context directly cmdCtx := ctx // This profile allows reading/writing to the working directory and /tmp, // but denies writing to other system locations by default (implicitly, though 'allow default' is permissive). // Use a basic profile for now. wrappedCommand := fmt.Sprintf("sandbox-exec -p %q /bin/bash -c %q", "(version 1) (allow default)", command) cmd := exec.CommandContext(cmdCtx, "/bin/bash", "-c", wrappedCommand) cmd.Dir = workDir cmd.Env = env var stdout, stderr bytes.Buffer cmd.Stdout = &stdout cmd.Stderr = &stderr err := cmd.Run() result := &ExecResult{ Command: command, Stdout: stdout.String(), Stderr: stderr.String(), ExitCode: 0, } if err != nil { result.Error = err.Error() if exitErr, ok := err.(*exec.ExitError); ok { result.ExitCode = exitErr.ExitCode() } else { result.ExitCode = 1 } } return result, nil } // Close is a no-op for Seatbelt executor. func (e *Seatbelt) Close(ctx context.Context) error { return nil } ================================================ FILE: pkg/sandbox/seatbelt_executor_others.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //go:build !darwin package sandbox import ( "context" "fmt" ) // Seatbelt executes commands in a seatbelt sandbox. type Seatbelt struct{} // NewSeatbeltExecutor creates a new SeatbeltExecutor. func NewSeatbeltExecutor() *Seatbelt { return &Seatbelt{} } // Execute executes the command in the seatbelt sandbox. func (e *Seatbelt) Execute(ctx context.Context, command string, env []string, workDir string) (*ExecResult, error) { return nil, fmt.Errorf("seatbelt sandbox is only supported on macOS") } // Close is a no-op for Seatbelt executor. func (e *Seatbelt) Close(ctx context.Context) error { return nil } ================================================ FILE: pkg/sessions/filesystem.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package sessions import ( "bufio" "encoding/json" "errors" "os" "path/filepath" "sort" "sync" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "sigs.k8s.io/yaml" ) type filesystemStore struct { basePath string } func newFilesystemStore(basePath string) Store { return &filesystemStore{basePath: basePath} } func (f *filesystemStore) GetSession(id string) (*api.Session, error) { sessionPath := filepath.Join(f.basePath, id) metadataPath := filepath.Join(sessionPath, "metadata.yaml") metadataBytes, err := os.ReadFile(metadataPath) if err != nil { if errors.Is(err, os.ErrNotExist) { return nil, errors.New("session not found") } return nil, err } var meta Metadata if err := yaml.Unmarshal(metadataBytes, &meta); err != nil { return nil, err } chatStore := NewFileChatMessageStore(sessionPath) return &api.Session{ ID: id, ProviderID: meta.ProviderID, ModelID: meta.ModelID, AgentState: api.AgentStateIdle, CreatedAt: meta.CreatedAt, LastModified: meta.LastAccessed, ChatMessageStore: chatStore, }, nil } func (f *filesystemStore) CreateSession(session *api.Session) error { sessionPath := filepath.Join(f.basePath, session.ID) if err := os.MkdirAll(sessionPath, 0o755); err != nil { return err } chatStore := NewFileChatMessageStore(sessionPath) session.ChatMessageStore = chatStore meta := Metadata{ ProviderID: session.ProviderID, ModelID: session.ModelID, CreatedAt: session.CreatedAt, LastAccessed: session.LastModified, } data, err := yaml.Marshal(meta) if err != nil { return err } return os.WriteFile(filepath.Join(sessionPath, "metadata.yaml"), data, 0o644) } func (f *filesystemStore) UpdateSession(session *api.Session) error { sessionPath := filepath.Join(f.basePath, session.ID) metadataPath := filepath.Join(sessionPath, "metadata.yaml") metadataBytes, err := os.ReadFile(metadataPath) if err != nil { if errors.Is(err, os.ErrNotExist) { return errors.New("session not found") } return err } var meta Metadata if err := yaml.Unmarshal(metadataBytes, &meta); err != nil { return err } meta.ProviderID = session.ProviderID meta.ModelID = session.ModelID meta.LastAccessed = session.LastModified data, err := yaml.Marshal(meta) if err != nil { return err } return os.WriteFile(metadataPath, data, 0o644) } func (f *filesystemStore) ListSessions() ([]*api.Session, error) { entries, err := os.ReadDir(f.basePath) if err != nil { if errors.Is(err, os.ErrNotExist) { return []*api.Session{}, nil } return nil, err } sessions := make([]*api.Session, 0, len(entries)) for _, entry := range entries { if !entry.IsDir() { continue } session, err := f.GetSession(entry.Name()) if err != nil { return nil, err } sessions = append(sessions, session) } sort.Slice(sessions, func(i, j int) bool { return sessions[i].LastModified.After(sessions[j].LastModified) }) return sessions, nil } func (f *filesystemStore) DeleteSession(id string) error { sessionPath := filepath.Join(f.basePath, id) return os.RemoveAll(sessionPath) } // FileChatMessageStore implements api.ChatMessageStore by persisting history to disk. type FileChatMessageStore struct { Path string mu sync.Mutex } // NewFileChatMessageStore creates a new file-backed chat message store. func NewFileChatMessageStore(path string) *FileChatMessageStore { return &FileChatMessageStore{Path: path} } // HistoryPath returns the location of the history file for this session. func (s *FileChatMessageStore) HistoryPath() string { return filepath.Join(s.Path, "history.json") } // AddChatMessage appends a message to the existing history on disk. func (s *FileChatMessageStore) AddChatMessage(record *api.Message) error { s.mu.Lock() defer s.mu.Unlock() // Ensure directory exists if err := os.MkdirAll(s.Path, 0o755); err != nil { return err } path := s.HistoryPath() // Check for legacy format and migrate if needed isLegacy := false if f, err := os.Open(path); err == nil { buf := make([]byte, 1) if _, err := f.Read(buf); err == nil && buf[0] == '[' { isLegacy = true } f.Close() } if isLegacy { // Read all messages (handles legacy format) messages, err := s.readMessages() if err != nil { return err } messages = append(messages, record) return s.writeMessages(messages) } // Normal append for JSONL or new files data, err := json.Marshal(record) if err != nil { return err } f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644) if err != nil { return err } defer f.Close() if _, err := f.Write(data); err != nil { return err } if _, err := f.WriteString("\n"); err != nil { return err } return nil } // SetChatMessages replaces the history file with the provided messages. func (s *FileChatMessageStore) SetChatMessages(newHistory []*api.Message) error { s.mu.Lock() defer s.mu.Unlock() return s.writeMessages(newHistory) } // ChatMessages returns all persisted chat messages. func (s *FileChatMessageStore) ChatMessages() []*api.Message { s.mu.Lock() defer s.mu.Unlock() messages, err := s.readMessages() if err != nil { return []*api.Message{} } return messages } // ClearChatMessages truncates the history file, leaving an empty array. func (s *FileChatMessageStore) ClearChatMessages() error { s.mu.Lock() defer s.mu.Unlock() return s.writeMessages([]*api.Message{}) } func (s *FileChatMessageStore) readMessages() ([]*api.Message, error) { path := s.HistoryPath() f, err := os.Open(path) if errors.Is(err, os.ErrNotExist) { return []*api.Message{}, nil } if err != nil { return nil, err } defer f.Close() // Check if the file is empty stat, err := f.Stat() if err != nil { return nil, err } if stat.Size() == 0 { return []*api.Message{}, nil } // Peek at the first byte to determine format // If it starts with '[', it's a legacy JSON array // Otherwise, assume JSONL buf := make([]byte, 1) if _, err := f.Read(buf); err != nil { return nil, err } // Reset file pointer if _, err := f.Seek(0, 0); err != nil { return nil, err } var messages []*api.Message if buf[0] == '[' { // Legacy JSON array format decoder := json.NewDecoder(f) if err := decoder.Decode(&messages); err != nil { return nil, err } return messages, nil } // JSONL format scanner := bufio.NewScanner(f) for scanner.Scan() { line := scanner.Bytes() if len(line) == 0 { continue } var msg api.Message if err := json.Unmarshal(line, &msg); err != nil { return nil, err } messages = append(messages, &msg) } if err := scanner.Err(); err != nil { return nil, err } return messages, nil } func (s *FileChatMessageStore) writeMessages(messages []*api.Message) error { if err := os.MkdirAll(s.Path, 0o755); err != nil { return err } f, err := os.OpenFile(s.HistoryPath(), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) if err != nil { return err } defer f.Close() for _, msg := range messages { data, err := json.Marshal(msg) if err != nil { return err } if _, err := f.Write(data); err != nil { return err } if _, err := f.WriteString("\n"); err != nil { return err } } return nil } ================================================ FILE: pkg/sessions/manager.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package sessions import ( "fmt" "math/rand" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) type SessionManager struct { store Store } func NewSessionManager(backend string) (*SessionManager, error) { var store Store var err error if backend == "" { // Try filesystem first store, err = NewStore("filesystem") if err != nil { // Fallback to memory store, err = NewStore("memory") } } else { store, err = NewStore(backend) } if err != nil { return nil, err } return &SessionManager{store: store}, nil } func (sm *SessionManager) NewSession(meta Metadata) (*api.Session, error) { suffix := fmt.Sprintf("%04d", rand.Intn(10000)) sessionID := time.Now().Format("20060102") + "-" + suffix now := time.Now() session := &api.Session{ ID: sessionID, Name: "Session " + sessionID, ProviderID: meta.ProviderID, ModelID: meta.ModelID, AgentState: api.AgentStateIdle, CreatedAt: now, LastModified: now, } if err := sm.store.CreateSession(session); err != nil { return nil, err } return session, nil } func (sm *SessionManager) ListSessions() ([]*api.Session, error) { return sm.store.ListSessions() } func (sm *SessionManager) FindSessionByID(id string) (*api.Session, error) { return sm.store.GetSession(id) } func (sm *SessionManager) DeleteSession(id string) error { return sm.store.DeleteSession(id) } func (sm *SessionManager) GetLatestSession() (*api.Session, error) { sessions, err := sm.store.ListSessions() if err != nil { return nil, err } if len(sessions) == 0 { return nil, nil } latest := sessions[0] for _, session := range sessions[1:] { if session.LastModified.After(latest.LastModified) { latest = session } } return latest, nil } func (sm *SessionManager) UpdateLastAccessed(session *api.Session) error { session.LastModified = time.Now() return sm.store.UpdateSession(session) } ================================================ FILE: pkg/sessions/memory.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package sessions import ( "errors" "sort" "sync" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) type memoryStore struct { mu sync.RWMutex sessions map[string]*api.Session } func newMemoryStore() Store { return &memoryStore{sessions: make(map[string]*api.Session)} } func (m *memoryStore) GetSession(id string) (*api.Session, error) { m.mu.RLock() defer m.mu.RUnlock() session, ok := m.sessions[id] if !ok { return nil, errors.New("session not found") } return session, nil } func (m *memoryStore) CreateSession(session *api.Session) error { m.mu.Lock() defer m.mu.Unlock() if _, exists := m.sessions[session.ID]; exists { return errors.New("session already exists") } if session.ChatMessageStore == nil { session.ChatMessageStore = NewInMemoryChatStore() } m.sessions[session.ID] = session return nil } func (m *memoryStore) UpdateSession(session *api.Session) error { m.mu.Lock() defer m.mu.Unlock() if _, exists := m.sessions[session.ID]; !exists { return errors.New("session not found") } m.sessions[session.ID] = session return nil } func (m *memoryStore) ListSessions() ([]*api.Session, error) { m.mu.RLock() defer m.mu.RUnlock() sessions := make([]*api.Session, 0, len(m.sessions)) for _, session := range m.sessions { sessions = append(sessions, session) } sort.Slice(sessions, func(i, j int) bool { return sessions[i].LastModified.After(sessions[j].LastModified) }) return sessions, nil } func (m *memoryStore) DeleteSession(id string) error { m.mu.Lock() defer m.mu.Unlock() if _, exists := m.sessions[id]; !exists { return errors.New("session not found") } delete(m.sessions, id) return nil } // InMemoryChatStore is an in-memory implementation of the api.ChatMessageStore interface. // It stores chat messages in a slice and is safe for concurrent use. type InMemoryChatStore struct { mu sync.RWMutex messages []*api.Message } // NewInMemoryChatStore creates a new InMemoryChatStore. func NewInMemoryChatStore() *InMemoryChatStore { return &InMemoryChatStore{ messages: make([]*api.Message, 0), } } // AddChatMessage adds a message to the store. func (s *InMemoryChatStore) AddChatMessage(record *api.Message) error { s.mu.Lock() defer s.mu.Unlock() s.messages = append(s.messages, record) return nil } // SetChatMessages replaces the entire chat history with a new one. func (s *InMemoryChatStore) SetChatMessages(newHistory []*api.Message) error { s.mu.Lock() defer s.mu.Unlock() s.messages = newHistory return nil } // ChatMessages returns all chat messages from the store. func (s *InMemoryChatStore) ChatMessages() []*api.Message { s.mu.RLock() defer s.mu.RUnlock() messageCopy := make([]*api.Message, len(s.messages)) copy(messageCopy, s.messages) return messageCopy } // ClearChatMessages removes all messages from the store. func (s *InMemoryChatStore) ClearChatMessages() error { s.mu.Lock() defer s.mu.Unlock() s.messages = make([]*api.Message, 0) return nil } ================================================ FILE: pkg/sessions/store.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package sessions import ( "fmt" "os" "path/filepath" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" ) const sessionsDirName = "sessions" type Metadata struct { ProviderID string `json:"providerID"` ModelID string `json:"modelID"` CreatedAt time.Time `json:"createdAt"` LastAccessed time.Time `json:"lastAccessed"` } var defaultMemoryStore Store = newMemoryStore() type Store interface { GetSession(id string) (*api.Session, error) CreateSession(session *api.Session) error UpdateSession(session *api.Session) error ListSessions() ([]*api.Session, error) DeleteSession(id string) error } func NewStore(backend string) (Store, error) { switch backend { case "memory": return defaultMemoryStore, nil case "filesystem": basePath, err := defaultFilesystemBasePath() if err != nil { return nil, err } if err := os.MkdirAll(basePath, 0o755); err != nil { return nil, err } return newFilesystemStore(basePath), nil default: return nil, fmt.Errorf("unsupported sessions backend: %s", backend) } } func defaultFilesystemBasePath() (string, error) { home, err := os.UserHomeDir() if err != nil { return "", err } return filepath.Join(home, ".kubectl-ai", sessionsDirName), nil } ================================================ FILE: pkg/tools/bash_tool.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "fmt" "os" "path/filepath" "runtime" "strings" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" ) const ( defaultBashBin = "/bin/bash" ) // expandShellVar expands shell variables and syntax using bash func expandShellVar(value string) (string, error) { if strings.Contains(value, "~") { if len(value) >= 2 && value[0] == '~' && os.IsPathSeparator(value[1]) { if runtime.GOOS == "windows" { value = filepath.Join(os.Getenv("USERPROFILE"), value[2:]) } else { value = filepath.Join(os.Getenv("HOME"), value[2:]) } } } return os.ExpandEnv(value), nil } type BashTool struct { executor sandbox.Executor } func NewBashTool(executor sandbox.Executor) *BashTool { return &BashTool{executor: executor} } func (t *BashTool) Name() string { return "bash" } func (t *BashTool) Description() string { return "Executes a bash command. Use this tool only when you need to execute a shell command." } func (t *BashTool) FunctionDefinition() *gollm.FunctionDefinition { return &gollm.FunctionDefinition{ Name: t.Name(), Description: t.Description(), Parameters: &gollm.Schema{ Type: gollm.TypeObject, Properties: map[string]*gollm.Schema{ "command": { Type: gollm.TypeString, Description: `The bash command to execute.`, }, "modifies_resource": { Type: gollm.TypeString, Description: `Whether the command modifies a kubernetes resource. Possible values: - "yes" if the command modifies a resource - "no" if the command does not modify a resource - "unknown" if the command's effect on the resource is unknown `, }, }, }, } } func (t *BashTool) Run(ctx context.Context, args map[string]any) (any, error) { kubeconfig := ctx.Value(KubeconfigKey).(string) workDir := ctx.Value(WorkDirKey).(string) command := args["command"].(string) if err := validateCommand(command); err != nil { return &sandbox.ExecResult{Command: command, Error: err.Error()}, nil } // Prepare environment env := os.Environ() if kubeconfig != "" { kubeconfig, err := expandShellVar(kubeconfig) if err != nil { return nil, err } env = append(env, "KUBECONFIG="+kubeconfig) } return ExecuteWithStreamingHandling(ctx, t.executor, command, workDir, env, DetectKubectlStreaming) } func validateCommand(command string) error { if strings.Contains(command, "kubectl edit") { return fmt.Errorf("interactive mode not supported for kubectl, please use non-interactive commands") } if strings.Contains(command, "kubectl port-forward") { return fmt.Errorf("port-forwarding is not allowed because assistant is running in an unattended mode, please try some other alternative") } return nil } func (t *BashTool) IsInteractive(args map[string]any) (bool, error) { commandVal, ok := args["command"] if !ok || commandVal == nil { return false, nil } command, ok := commandVal.(string) if !ok { return false, nil } return IsInteractiveCommand(command) } // CheckModifiesResource determines if the command modifies kubernetes resources // This is used for permission checks before command execution // Returns "yes", "no", or "unknown" func (t *BashTool) CheckModifiesResource(args map[string]any) string { command, ok := args["command"].(string) if !ok { return "unknown" } if strings.Contains(command, "kubectl") { return kubectlModifiesResource(command) } return "unknown" } ================================================ FILE: pkg/tools/custom_tool.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "fmt" "os" "strings" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" "mvdan.cc/sh/v3/syntax" ) // CustomToolConfig defines the structure for configuring a custom tool. type CustomToolConfig struct { Name string `yaml:"name"` Description string `yaml:"description"` Command string `yaml:"command"` CommandDesc string `yaml:"command_desc"` IsInteractive bool `yaml:"is_interactive"` } // CustomTool implements the Tool interface for external commands. type CustomTool struct { config CustomToolConfig executor sandbox.Executor } // NewCustomTool creates a new CustomTool instance. func NewCustomTool(config CustomToolConfig) (*CustomTool, error) { if config.Name == "" { return nil, fmt.Errorf("custom tool name cannot be empty") } if len(config.Command) == 0 { return nil, fmt.Errorf("custom tool command cannot be empty for tool %q", config.Name) } return &CustomTool{config: config}, nil } // Name returns the tool's name. func (t *CustomTool) Name() string { return t.config.Name } // Description returns the tool's description from its function definition. func (t *CustomTool) Description() string { return t.config.Description } // FunctionDefinition returns the tool's function definition. func (t *CustomTool) FunctionDefinition() *gollm.FunctionDefinition { return &gollm.FunctionDefinition{ Name: t.Name(), Description: t.Description(), Parameters: &gollm.Schema{ Type: gollm.TypeObject, Properties: map[string]*gollm.Schema{ "command": { Type: gollm.TypeString, Description: t.config.CommandDesc, }, "modifies_resource": { Type: gollm.TypeString, Description: `Whether the command modifies a resource. Possible values: - "yes" if the command modifies a resource - "no" if the command does not modify a resource - "unknown" if the command's effect on the resource is unknown `, }, }, }, } } // addCommandPrefix adds the tool's command prefix to the input command if needed. // It only adds the prefix if the command is a simple command (no pipes, etc.) // and doesn't already start with the prefix. // TODO(droot): This will not be needed when models improve on following instructions // and specify the complete command to execute. func (t *CustomTool) addCommandPrefix(inputCmd string) (string, error) { // Parse the command to check if it's a simple command parser := syntax.NewParser() prog, err := parser.Parse(strings.NewReader(inputCmd), "") if err != nil { return "", fmt.Errorf("failed to parse command: %w", err) } // Check if it's a simple command (no pipes, redirects, etc.) if len(prog.Stmts) != 1 { return inputCmd, nil } stmt := prog.Stmts[0] if stmt.Background || stmt.Coprocess || stmt.Negated || len(stmt.Redirs) > 0 { return inputCmd, nil } // Check if it's a simple call expression if _, ok := stmt.Cmd.(*syntax.CallExpr); !ok { return inputCmd, nil } // If we get here, it's a simple command without the prefix if strings.HasPrefix(inputCmd, t.config.Command) { return inputCmd, nil } return t.config.Command + " " + inputCmd, nil } // Run executes the external command defined for the custom tool. func (t *CustomTool) Run(ctx context.Context, args map[string]any) (any, error) { var command string cmdVal, ok := args["command"] if !ok { return nil, fmt.Errorf("command not found in args") } command = cmdVal.(string) command, err := t.addCommandPrefix(command) if err != nil { return nil, fmt.Errorf("failed to process command: %w", err) } workDir := ctx.Value(WorkDirKey).(string) env := os.Environ() // Use the injected executor, or fallback to local if not set (e.g. for global instance) executor := t.executor if executor == nil { executor = sandbox.NewLocalExecutor() } // Execute the command return executor.Execute(ctx, command, env, workDir) } // CheckModifiesResource determines if the command modifies resources // For custom tools, we'll conservatively assume they might modify resources // unless we have specific knowledge otherwise // Returns "yes", "no", or "unknown" func (t *CustomTool) CheckModifiesResource(args map[string]any) string { // For custom tools, we'll conservatively use "unknown" since we can't return "unknown" } // CloneWithExecutor creates a copy of the CustomTool with the given executor. // This is used to create a session-specific instance of the tool. func (t *CustomTool) CloneWithExecutor(executor sandbox.Executor) *CustomTool { return &CustomTool{ config: t.config, executor: executor, } } ================================================ FILE: pkg/tools/custom_tool_test.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "strings" "testing" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" ) func TestCustomTool_AddCommandPrefix(t *testing.T) { tests := []struct { name string configCommand string inputCommand string expectedOutput string expectError bool }{ { name: "simple command without prefix", configCommand: "gcloud", inputCommand: "compute instances list", expectedOutput: "gcloud compute instances list", expectError: false, }, { name: "simple command with prefix", configCommand: "gcloud", inputCommand: "gcloud compute instances list", expectedOutput: "gcloud compute instances list", expectError: false, }, { name: "command with pipe", configCommand: "gcloud", inputCommand: "compute instances list | grep test", expectedOutput: "compute instances list | grep test", expectError: false, }, { name: "command with redirect", configCommand: "gcloud", inputCommand: "compute instances list > instances.txt", expectedOutput: "compute instances list > instances.txt", expectError: false, }, { name: "command with background", configCommand: "gcloud", inputCommand: "compute instances list &", expectedOutput: "compute instances list &", expectError: false, }, { name: "command with subshell", configCommand: "gcloud", inputCommand: "(compute instances list)", expectedOutput: "(compute instances list)", expectError: false, }, { name: "command with multiple statements", configCommand: "gcloud", inputCommand: "compute instances list; compute disks list", expectedOutput: "compute instances list; compute disks list", expectError: false, }, { name: "invalid shell syntax", configCommand: "gcloud", inputCommand: "compute instances list |", expectedOutput: "", expectError: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tool := &CustomTool{ config: CustomToolConfig{ Command: tt.configCommand, }, } output, err := tool.addCommandPrefix(tt.inputCommand) if tt.expectError { if err == nil { t.Errorf("expected error but got none") } return } if err != nil { t.Errorf("unexpected error: %v", err) return } if output != tt.expectedOutput { t.Errorf("expected %q, got %q", tt.expectedOutput, output) } }) } } // MockExecutor implements sandbox.Executor for testing type MockExecutor struct { CapturedCommand string CapturedEnv []string CapturedWorkDir string } func (m *MockExecutor) Execute(ctx context.Context, command string, env []string, workDir string) (*sandbox.ExecResult, error) { m.CapturedCommand = command m.CapturedEnv = env m.CapturedWorkDir = workDir return &sandbox.ExecResult{Stdout: "executed"}, nil } func (m *MockExecutor) Close(ctx context.Context) error { return nil } func TestCustomTool_CloneWithExecutor(t *testing.T) { config := CustomToolConfig{ Name: "test-tool", Command: "echo", } tool, err := NewCustomTool(config) if err != nil { t.Fatalf("failed to create tool: %v", err) } mockExec := &MockExecutor{} clonedTool := tool.CloneWithExecutor(mockExec) ctx := context.WithValue(context.Background(), WorkDirKey, "/tmp") args := map[string]any{ "command": "hello", } result, err := clonedTool.Run(ctx, args) if err != nil { t.Fatalf("tool run failed: %v", err) } execResult, ok := result.(*sandbox.ExecResult) if !ok { t.Fatalf("expected *sandbox.ExecResult, got %T", result) } if execResult.Stdout != "executed" { t.Errorf("expected Stdout 'executed', got %q", execResult.Stdout) } if mockExec.CapturedCommand != "echo hello" { t.Errorf("expected command 'echo hello', got %q", mockExec.CapturedCommand) } if !strings.Contains(strings.Join(mockExec.CapturedEnv, "\n"), "PATH") { t.Errorf("expected captured environment to contain PATH") } if mockExec.CapturedWorkDir != "/tmp" { t.Errorf("expected workdir '/tmp', got %q", mockExec.CapturedWorkDir) } } ================================================ FILE: pkg/tools/interfaces.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" ) type Tool interface { // Name is the identifier for the tool; we pass this to the LLM. // The LLM uses this name when it wants to invoke the tool. // It should be meaningful and (we think) camel_case as (we think) that works better with most LLMs. Name() string // Description is an additional description that gives the LLM instructions on when to use the tool. Description() string // FunctionDefinition provides the full schema for the parameters to be used when invoking the tool. // The Description fields provides hints that the LLM may use to use the tool more effectively/correctly. FunctionDefinition() *gollm.FunctionDefinition // Run invokes the tool, the agent calls this when the LLM requests tool invocation. Run(ctx context.Context, args map[string]any) (any, error) // IsInteractive checks if a command is interactive // If the command is interactive, we need to handle it differently in the agent // Returns true if interactive, with an error explaining why it's interactive IsInteractive(args map[string]any) (bool, error) // CheckModifiesResource determines if the command modifies resources // This is used for permission checks before command execution // Returns "yes", "no", or "unknown" CheckModifiesResource(args map[string]any) string } ================================================ FILE: pkg/tools/kubectl_filter.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "strings" "k8s.io/klog/v2" "mvdan.cc/sh/v3/syntax" ) // Package-level constants for kubectl operations var ( readOnlyOps = map[string]bool{ "get": true, "describe": true, "explain": true, "top": true, "logs": true, "api-resources": true, "api-versions": true, "version": true, "config": true, "cluster-info": true, "wait": true, "auth": true, "diff": true, "kustomize": true, "help": true, "options": true, "proxy": true, "completion": true, "convert": true, "events": true, "port-forward": true, "can-i": true, "whoami": true, } writeOps = map[string]bool{ "create": true, "apply": true, "edit": true, "delete": true, "patch": true, "replace": true, "scale": true, "autoscale": true, "expose": true, "run": true, "exec": true, "set": true, "label": true, "annotate": true, "taint": true, "drain": true, "cordon": true, "uncordon": true, "debug": true, "attach": true, "cp": true, "reconcile": true, "approve": true, "deny": true, "certificate": true, } readOnlySubOps = map[string]map[string]bool{ "rollout": { "history": true, "status": true, }, } writeSubOps = map[string]map[string]bool{ "rollout": { "pause": true, "restart": true, "resume": true, "undo": true, }, } ) // KubectlModifiesResource analyzes a kubectl command to determine if it modifies resources func kubectlModifiesResource(command string) string { parser := syntax.NewParser() file, err := parser.Parse(strings.NewReader(command), "") if err != nil { klog.Errorf("Failed to parse kubectl command: %v, command: %q", err, command) return "unknown" } hasReadCommand := false foundWrite := false numCmds := 0 // Single pass through all command calls syntax.Walk(file, func(node syntax.Node) bool { if call, ok := node.(*syntax.CallExpr); ok { result := analyzeCall(call) // If we find any write operation, mark it and stop if result == "yes" { foundWrite = true return false // Stop walking } // Track if we found any read operations if result == "no" { hasReadCommand = true } numCmds++ if numCmds > 1 { return false // Stop walking if more then one command is found } } return true }) if numCmds > 1 { // if it's a composite bash command, we should err on the side of caution and return unknown // to prevent exfilteration attacks https://simonwillison.net/2025/Jun/16/the-lethal-trifecta/ klog.Infof("KubectlModifiesResource result: unknown for command: %q, multiple commands (%d) found", command, numCmds) return "unknown" } // Return results based on what we found if foundWrite { klog.Infof("KubectlModifiesResource result: yes (write operation found) for command: %q", command) return "yes" } if hasReadCommand { klog.Infof("KubectlModifiesResource result: no (read-only) for command: %q", command) return "no" } // Default to unknown if no recognized kubectl commands found klog.Infof("KubectlModifiesResource result: unknown for command: %q", command) return "unknown" } func analyzeCall(call *syntax.CallExpr) string { if call == nil || len(call.Args) == 0 { klog.Warning("analyzeCall: call is nil or has no args") return "unknown" } // Extract command and arguments var args []string for _, arg := range call.Args { lit := arg.Lit() if lit == "" { var sb strings.Builder syntax.NewPrinter().Print(&sb, arg) lit = strings.Trim(sb.String(), "'\"") } if lit != "" { args = append(args, lit) } } if len(args) == 0 { klog.Warning("analyzeCall: no arguments extracted from call") return "unknown" } // Check if first argument is kubectl firstArg := args[0] // Reject quoted arguments (e.g., '"/path/kubectl"') if (strings.HasPrefix(firstArg, "'") && strings.HasSuffix(firstArg, "'")) || (strings.HasPrefix(firstArg, "\"") && strings.HasSuffix(firstArg, "\"")) { klog.V(2).Infof("analyzeCall: first arg is quoted: %q", firstArg) return "unknown" } // Check if this is kubectl if !strings.Contains(firstArg, "kubectl") { klog.V(2).Infof("analyzeCall: first arg does not contain kubectl: %q", firstArg) return "unknown" } klog.V(2).Infof("analyzeCall: found kubectl: %q", firstArg) // Check for boolean or spaced key-value flags before the verb for _, arg := range args[1:] { if !strings.HasPrefix(arg, "-") { break } // If flag does not contain '=', it's boolean or spaced key-value if !strings.Contains(arg, "=") { klog.Warningf("analyzeCall: boolean or spaced key-value flag before verb: %q", arg) return "unknown" } } // Parse kubectl arguments to extract verb, subverb, and flags verb, subVerb, hasDryRun := parseKubectlArgs(args[1:]) if verb == "" { klog.Warningf("analyzeCall: no verb found after kubectl in args: %v", args) return "unknown" } // Check standard operations - write operations first (prioritize immediate detection) if (writeOps[verb] || writeSubOps[verb][subVerb]) && !hasDryRun { klog.V(1).Infof("analyzeCall: write op for verb=%q subVerb=%q", verb, subVerb) return "yes" } // Check read-only operations or dry-run write operations if (readOnlyOps[verb] || readOnlySubOps[verb][subVerb]) || ((writeOps[verb] || writeSubOps[verb][subVerb]) && hasDryRun) { klog.V(1).Infof("analyzeCall: read op for verb=%q subVerb=%q (dry-run=%v)", verb, subVerb, hasDryRun) return "no" } klog.V(1).Infof("analyzeCall: unknown op for verb=%q subVerb=%q", verb, subVerb) return "unknown" } // parseKubectlArgs extracts verb, subverb, and dry-run flag from kubectl arguments func parseKubectlArgs(args []string) (verb, subVerb string, hasDryRun bool) { for _, arg := range args { if strings.HasPrefix(arg, "--dry-run") { hasDryRun = true } if !strings.HasPrefix(arg, "-") { if verb == "" { verb = arg } else if subVerb == "" { subVerb = arg } } } return verb, subVerb, hasDryRun } ================================================ FILE: pkg/tools/kubectl_filter_test.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "path/filepath" "strings" "testing" "mvdan.cc/sh/v3/syntax" ) func TestKubectlModifiesResource(t *testing.T) { // Group test cases by category testCases := map[string][]struct { name string command string expected string }{ "read-only commands": { {"Get pods", "kubectl get pods", "no"}, {"Describe pod", "kubectl describe pod nginx", "no"}, {"Port-forward", "kubectl port-forward pod/nginx 8080:80", "no"}, {"Port-forward with service", "kubectl port-forward svc/nginx 8080:80", "no"}, {"Port-forward complex", "kubectl port-forward deployment/myapp 8080:8080 9000:9000", "no"}, {"Port-forward background", "kubectl port-forward svc/nginx 8080:80 &", "no"}, {"Get with output", "kubectl get pods -o yaml", "no"}, {"Get with output redirection", "kubectl get pods > pods.txt", "no"}, {"Get with name", "kubectl get pod nginx", "no"}, {"Config view", "kubectl config view", "no"}, {"Version", "kubectl version", "no"}, {"API resources", "kubectl api-resources", "no"}, {"Explain", "kubectl explain pods", "no"}, {"Logs", "kubectl logs nginx", "no"}, {"Logs with follow", "kubectl logs nginx -f", "no"}, {"Watch pods", "kubectl get pods --watch", "no"}, {"Watch pods short", "kubectl get pods -w", "no"}, {"Rollout status", "kubectl rollout status deployment nginx", "no"}, {"Diff", "kubectl diff -f deployment.yaml", "no"}, {"Can-i", "kubectl auth can-i create pods", "no"}, {"Kustomize", "kubectl kustomize ./", "no"}, {"Convert", "kubectl convert -f pod.yaml --output-version=v1", "no"}, {"Events", "kubectl events", "no"}, {"Alpha debug", "kubectl alpha debug pod/nginx", "unknown"}, {"Auth whoami", "kubectl auth whoami", "no"}, }, "modifying commands": { {"Create pod", "kubectl create -f pod.yaml", "yes"}, {"Apply deployment", "kubectl apply -f deployment.yaml", "yes"}, {"Delete pod", "kubectl delete pod nginx", "yes"}, {"Scale deployment", "kubectl scale deployment nginx --replicas=3", "yes"}, {"Edit deployment", "kubectl edit deployment nginx", "yes"}, {"Patch service", "kubectl patch svc nginx -p '{\"spec\":{\"type\":\"NodePort\"}}'", "yes"}, {"Label pod", "kubectl label pod nginx app=web", "yes"}, {"Annotate", "kubectl annotate pods nginx description='my nginx'", "yes"}, {"Rollout restart", "kubectl rollout restart deployment nginx", "yes"}, {"Set image", "kubectl set image deployment/nginx nginx=nginx:latest", "yes"}, {"Exec into pod", "kubectl exec -n demo tgi-pod -- nvidia-smi", "yes"}, {"Taint node", "kubectl taint nodes node1 key=value:NoSchedule", "yes"}, {"Run pod", "kubectl run nginx --image=nginx", "yes"}, {"Config set-context", "kubectl config set-context my-context", "no"}, {"Exec command", "kubectl exec nginx -- rm -rf /", "yes"}, {"Cordon node", "kubectl cordon node1", "yes"}, {"Uncordon node", "kubectl uncordon node1", "yes"}, {"Drain node", "kubectl drain node1", "yes"}, {"Certificate approve", "kubectl certificate approve csr-12345", "yes"}, }, "special cases": { {"Dry run create", "kubectl create -f pod.yaml --dry-run=client", "no"}, {"Dry run apply", "kubectl apply -f deployment.yaml --dry-run", "no"}, {"Apply with server dry-run", "kubectl apply -f pod.yaml --dry-run=server", "no"}, {"Delete with dry-run", "kubectl delete pod nginx --dry-run client", "no"}, }, "edge cases": { {"Command with pipe", "kubectl get pods | grep nginx", "unknown"}, {"Command with backticks", "kubectl get pod `cat podname.txt`", "unknown"}, {"Complex path", "\"/path with spaces/kubectl\" get pods", "no"}, {"Command with env var", "KUBECONFIG=/path/to/config kubectl get pods", "no"}, {"Not kubectl command", "ls -la", "unknown"}, {"Multiple spaces", "kubectl get pods", "no"}, {"Complex command with variables", "kubectl get pods -l app=$APP_NAME -n $NAMESPACE", "no"}, {"Command with quotes", "kubectl get pods -l \"app=my app\"", "no"}, {"Command with escaped quotes", "kubectl patch configmap my-config --patch \"{\\\"data\\\":{\\\"key\\\":\\\"new-value\\\"}}\"", "yes"}, {"Complex env vars", "KUBECONFIG=/path/to/config NS=default kubectl get pods -n $NS", "no"}, {"Command with multiple env vars", "KUBECONFIG=/config KUBECTL_EXTERNAL_DIFF=\"diff -u\" kubectl diff -f file.yaml", "no"}, {"Sequential commands with semicolon", "kubectl get ns; kubectl create ns test", "yes"}, {"Multiple safe commands", "kubectl get pods; kubectl get deployments", "unknown"}, {"Mix safe and unsafe with result", "kubectl get pods && kubectl delete pod bad-pod", "yes"}, {"Mix with initial unsafe", "kubectl delete pod bad-pod && kubectl get pods", "yes"}, {"Kubectl alias k", "k get pods", "unknown"}, {"Full path with arguments", "/usr/local/custom/kubectl --kubeconfig=/path/config get pods", "no"}, {"Complex jsonpath", "kubectl get pods -o=jsonpath='{range .items[*]}{.metadata.name}{\"\\t\"}{.status.phase}{\"\\n\"}{end}'", "no"}, {"Custom columns", "kubectl get pods -o=custom-columns=NAME:.metadata.name,STATUS:.status.phase", "no"}, {"Impersonation", "kubectl get pods --as=system:serviceaccount:default:deployer", "no"}, {"With token", "kubectl --token=eyJhbGciOiJSUzI1NiIsImtpZCI6IiJ9... get pods", "no"}, {"Weird spacing", "kubectl get pods -n default", "no"}, {"kubectl in shell script with line continuation", "kubectl get pods \\\n --namespace=production", "no"}, {"kubectl command split across lines in CI/CD", "kubectl delete pod \\\n my-pod-name \\\n --grace-period=30", "yes"}, {"kubectl with multiline YAML pipe", "echo 'apiVersion: v1\nkind: Pod' | kubectl apply -f -", "yes"}, {"kubectl logs with line breaks in shell", "kubectl logs \\\n deployment/my-app \\\n --follow", "no"}, {"Watch with selector", "kubectl get pods --selector app=nginx --watch", "no"}, {"Negative watch timeout", "kubectl get pods --watch-only --timeout=10s", "no"}, {"Flags after name", "kubectl delete pod mypod --now --grace-period=0", "yes"}, {"Server-side apply", "kubectl apply -f deploy.yaml --server-side", "yes"}, {"Field manager", "kubectl apply -f deploy.yaml --field-manager=controller", "yes"}, {"Create service account", "kubectl create serviceaccount jenkins", "yes"}, {"Create role binding", "kubectl create rolebinding admin --clusterrole=admin --user=user1 --namespace=default", "yes"}, {"Versioned kubectl", "kubectl.1.24 get pods", "no"}, {"Config set credentials", "kubectl config set-credentials cluster-admin --token=secret", "no"}, {"Config view with flatten", "kubectl config view --flatten", "no"}, {"Config view with output", "kubectl config view -o json", "no"}, {"Config use-context", "kubectl config use-context production", "no"}, {"Label with special characters", "kubectl label pod nginx 'app.kubernetes.io/name=nginx-controller'", "yes"}, {"Jsonpath with quotes", "kubectl get pods -o jsonpath='{.items[0].metadata.name}'", "no"}, {"Command with grep", "kubectl get pods | grep -v Completed", "unknown"}, {"Command with awk", "kubectl get pods | awk '{print $1}'", "unknown"}, {"Delete with force", "kubectl delete pod stuck-pod --force --grace-period=0", "yes"}, {"Custom resource get", "kubectl get virtualmachines", "no"}, {"Custom resource apply", "kubectl apply -f vm-instance.yaml", "yes"}, {"Multiple input files", "kubectl delete -f file1.yaml -f file2.yaml", "yes"}, {"URL as input file", "kubectl apply -f https://example.com/manifest.yaml", "yes"}, {"Input from stdin", "cat file.yaml | kubectl apply -f -", "yes"}, {"Proxy command", "kubectl proxy --port=8080", "no"}, {"Attach command", "kubectl attach mypod -i", "yes"}, {"Copy files", "kubectl cp mypod:/tmp/foo /tmp/bar", "yes"}, {"Rollout status with flags", "kubectl rollout --recursive=false status --timeout=0s deployment -w nginx", "no"}, }, } for category, cases := range testCases { t.Run(category, func(t *testing.T) { for _, tt := range cases { t.Run(tt.name, func(t *testing.T) { result := kubectlModifiesResource(tt.command) if result != tt.expected { t.Errorf("KubectlModifiesResource(%q) = %q, want %q", tt.command, result, tt.expected) } }) } }) } } // TestKubectlAnalyzerComponents tests the internal helper functions used by KubectlModifiesResource func TestKubectlAnalyzerComponents(t *testing.T) { t.Run("parseKubectlArgs detection", func(t *testing.T) { tests := []struct { command string verbExpected string subverbExpected string hasDryRunExpected bool }{ {"kubectl apply -f deploy.yaml --dry-run=client", "apply", "deploy.yaml", true}, {"kubectl apply -f deploy.yaml --dry-run", "apply", "deploy.yaml", true}, {"kubectl delete pod nginx --dry-run client", "delete", "pod", true}, {"kubectl delete pod nginx --dry-run=server", "delete", "pod", true}, {"kubectl apply -f deploy.yaml", "apply", "deploy.yaml", false}, {"kubectl get pods --dry", "get", "pods", false}, // Not a valid dry-run flag {"echo --dry-run", "", "", true}, // The current implementation doesn't check if it's kubectl {"kubectl rollout status deployment nginx", "rollout", "status", false}, } for _, tt := range tests { verb, subVerb, hasDryRun := parseKubectlArgs(strings.Split(tt.command, " ")[1:]) // Skip the first arg (kubectl) if verb != tt.verbExpected { t.Errorf("parseKubectlArgs(%q) verb = %q, want %q", tt.command, verb, tt.verbExpected) } if subVerb != tt.subverbExpected { t.Errorf("parseKubectlArgs(%q) subVerb = %q, want %q", tt.command, subVerb, tt.subverbExpected) } if hasDryRun != tt.hasDryRunExpected { t.Errorf("parseKubectlArgs(%q) hasDryRun = %v, want %v", tt.command, hasDryRun, tt.hasDryRunExpected) } } }) t.Run("command parsing", func(t *testing.T) { tests := []struct { command string expectedRes string }{ {"kubectl get pods", "no"}, {"kubectl apply -f deploy.yaml", "yes"}, {"ls -la", "unknown"}, // Not kubectl {"kubectl", "unknown"}, // Incomplete command {"kubectl; ls", "unknown"}, // Multiple commands } for _, tt := range tests { result := kubectlModifiesResource(tt.command) if result != tt.expectedRes { t.Errorf("KubectlModifiesResource(%q) = %q, want %q", tt.command, result, tt.expectedRes) } } }) } // TestKubectlCommandParsing tests kubectl command parsing focusing on realistic scenarios func TestKubectlCommandParsing(t *testing.T) { tests := []struct { name string command string expected string desc string }{ // Basic kubectl detection {"literal kubectl", "kubectl get pods", "no", "standard kubectl command"}, {"kubectl.exe", "kubectl.exe get pods", "no", "Windows executable"}, {"Unix path", "/usr/bin/kubectl get pods", "no", "Full Unix path"}, {"relative path", "./kubectl get services", "no", "relative path"}, {"nested path", "../bin/kubectl describe pod nginx", "no", "nested relative path"}, // Windows paths with forward slashes (works cross-platform) {"Windows forward slash", "C:/tools/kubectl.exe delete pod nginx", "yes", "Windows path with forward slash"}, // macOS/Homebrew paths {"macOS Homebrew", "/opt/homebrew/bin/kubectl get nodes", "no", "macOS Homebrew path"}, {"macOS Intel Homebrew", "/usr/local/bin/kubectl create namespace test", "yes", "macOS Intel Homebrew path"}, {"macOS Applications", "/Applications/Docker.app/Contents/Resources/bin/kubectl get all", "no", "macOS Docker Desktop kubectl"}, // Non-kubectl commands {"not kubectl", "k get pods", "unknown", "kubectl alias"}, {"kubectl suffix", "kubectl-1.28 get pods", "no", "kubectl with version suffix"}, {"kubectl prefix", "kubectl-proxy --port=8080", "unknown", "kubectl with additional suffix"}, {"different command", "kubectx production", "unknown", "different k8s tool"}, // Environment variables {"env var prefix", "KUBECONFIG=/path/config kubectl get pods", "no", "environment variable prefix"}, {"multiple env vars", "KUBECONFIG=/config NS=default kubectl apply -f deploy.yaml --dry-run", "no", "multiple environment variables"}, // Complex scenarios {"long path", "/very/long/path/to/kubectl get pods", "no", "very long path"}, {"flags before verb", "kubectl --context=prod --namespace=app get pods", "no", "global flags before verb"}, {"flags before verb mutating", "kubectl --replicas=3 scale deployment/nginx-deployment", "yes", "global flags before verb mutating"}, {"flags before verb without equals", "kubectl --context prod --namespace app get pods", "unknown", "global flags before verb without equals"}, {"no verb", "kubectl --help", "unknown", "kubectl with only flags"}, {"boolean flag before verb", "kubectl --verbose get pods", "unknown", "boolean flag before verb"}, {"boolean flag before verb mutating", "kubectl --force delete pod nginx", "unknown", "boolean flag before verb mutating"}, {"mixed flags before verb", "kubectl --context=prod --namespace app get pods", "unknown", "mixed non-spaced and spaced flags before verb"}, {"non-spaced key-value before verb non-mutating", "kubectl --namespace=default get pods", "no", "non-spaced key-value before verb non-mutating"}, {"non-spaced key-value before verb mutating", "kubectl --namespace=default delete pod nginx", "yes", "non-spaced key-value before verb mutating"}, {"flag after verb spaced", "kubectl get pods --context prod", "no", "spaced key-value flag after verb"}, {"flag after verb boolean", "kubectl get pods --verbose", "no", "boolean flag after verb"}, {"flag after verb mutating", "kubectl delete pod nginx --force", "yes", "boolean flag after verb mutating"}, {"flag with equals empty value before verb", "kubectl --namespace= get pods", "no", "non-spaced key-value with empty value before verb"}, {"unexpected arg before verb", "kubectl something get pods", "unknown", "unexpected arg before verb"}, {"multiple boolean flags before verb", "kubectl --verbose --debug get pods", "unknown", "multiple boolean flags before verb"}, {"multiple spaced flags before verb", "kubectl --context prod --namespace app get pods", "unknown", "multiple spaced flags before verb"}, {"multiple non-spaced flags before verb mutating", "kubectl --namespace=default --force=true delete pod nginx", "yes", "multiple non-spaced flags before verb mutating"}, {"multiple non-spaced flags before verb non-mutating", "kubectl --namespace=default --verbose=true get pods", "no", "multiple non-spaced flags before verb non-mutating"}, // Dry run scenarios {"dry run create", "/usr/bin/kubectl create -f pod.yaml --dry-run=client", "no", "dry run with path"}, {"dry run apply", "kubectl.exe apply -f deploy.yaml --dry-run", "no", "Windows dry run"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := kubectlModifiesResource(tt.command) if result != tt.expected { t.Errorf("KubectlModifiesResource(%q) = %q, want %q\nDescription: %s", tt.command, result, tt.expected, tt.desc) } }) } } // TestKubectlPathHandling tests the OS-agnostic path handling specifically func TestKubectlPathHandling(t *testing.T) { tests := []struct { name string binaryPath string shouldMatch bool description string }{ // Basic cases {"Standard kubectl", "kubectl", true, "Standard kubectl binary name"}, {"Windows kubectl.exe", "kubectl.exe", true, "Windows kubectl executable"}, {"Unix full path", "/usr/bin/kubectl", true, "Full Unix path to kubectl"}, {"Windows forward slash", "C:/tools/kubectl.exe", true, "Windows path with forward slashes"}, {"Relative path", "./kubectl", true, "Relative path to kubectl"}, {"macOS Homebrew", "/opt/homebrew/bin/kubectl", true, "macOS Homebrew path"}, // Non-kubectl binaries {"Not kubectl", "k", false, "Short alias should not match"}, {"kubectl with suffix", "kubectl-1.28", false, "kubectl with version suffix"}, {"kubectl prefix", "kubectl-proxy", false, "kubectl with additional suffix"}, {"Other binary", "kubectx", false, "Different binary altogether"}, {"Empty path", "", false, "Empty binary path"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Test the filepath.Base logic used in the actual function basename := filepath.Base(tt.binaryPath) isKubectl := (basename == "kubectl" || basename == "kubectl.exe") if isKubectl != tt.shouldMatch { t.Errorf("filepath.Base(%q) = %q, kubectl detection = %v, want %v\nDescription: %s", tt.binaryPath, basename, isKubectl, tt.shouldMatch, tt.description) } }) } } // TestKubectlDetectionLogic tests the core kubectl detection logic func TestKubectlDetectionLogic(t *testing.T) { // Simulate the kubectl detection logic from analyzeCall testKubectlDetection := func(arg string) bool { // Reject quoted arguments if (strings.HasPrefix(arg, "'") && strings.HasSuffix(arg, "'")) || (strings.HasPrefix(arg, "\"") && strings.HasSuffix(arg, "\"")) { return false } // Check if this is kubectl using OS-agnostic path handling basename := filepath.Base(arg) return basename == "kubectl" || basename == "kubectl.exe" } testCases := []struct { input string expected bool desc string }{ {"kubectl", true, "literal kubectl"}, {"kubectl.exe", true, "Windows executable"}, {"/usr/bin/kubectl", true, "Unix path"}, {"C:/tools/kubectl.exe", true, "Windows path with forward slash"}, {"./kubectl", true, "relative path"}, {"../bin/kubectl", true, "relative path with parent dir"}, {"/opt/homebrew/bin/kubectl", true, "macOS Homebrew path"}, {"'kubectl'", false, "quoted kubectl"}, {"\"/usr/bin/kubectl\"", false, "quoted path"}, {"not-kubectl", false, "different command"}, {"/usr/bin/k", false, "different command in path"}, {"kubectl-something", false, "kubectl with suffix"}, {"kubectl-1.28", false, "kubectl with version suffix"}, {"k", false, "kubectl alias"}, {"kubectx", false, "different k8s tool"}, {"", false, "empty string"}, } for _, tc := range testCases { t.Run(tc.desc, func(t *testing.T) { result := testKubectlDetection(tc.input) if result != tc.expected { t.Errorf("testKubectlDetection(%q) = %t, want %t (%s)", tc.input, result, tc.expected, tc.desc) } }) } } // TestShellParserBehavior tests how the shell parser handles different command structures // This helps us understand if we can simplify the kubectl detection logic func TestShellParserBehavior(t *testing.T) { testCommands := []struct { name string command string expected [][]string // expected args for each CallExpr }{ { name: "simple kubectl", command: "kubectl get pods", expected: [][]string{ {"kubectl", "get", "pods"}, }, }, { name: "kubectl with env var", command: "KUBECONFIG=/path/config kubectl get pods", expected: [][]string{ {"kubectl", "get", "pods"}, // env vars are handled separately }, }, { name: "sequential commands", command: "kubectl get pods; kubectl create pod", expected: [][]string{ {"kubectl", "get", "pods"}, {"kubectl", "create", "pod"}, }, }, { name: "kubectl with path", command: "/usr/bin/kubectl get pods", expected: [][]string{ {"/usr/bin/kubectl", "get", "pods"}, }, }, { name: "not kubectl", command: "ls -la", expected: [][]string{ {"ls", "-la"}, }, }, } parser := syntax.NewParser() for _, tt := range testCommands { t.Run(tt.name, func(t *testing.T) { file, err := parser.Parse(strings.NewReader(tt.command), "") if err != nil { t.Fatalf("Parse error for %q: %v", tt.command, err) } var actualCalls [][]string syntax.Walk(file, func(node syntax.Node) bool { if call, ok := node.(*syntax.CallExpr); ok { var args []string for _, arg := range call.Args { lit := arg.Lit() if lit == "" { var sb strings.Builder syntax.NewPrinter().Print(&sb, arg) lit = strings.Trim(sb.String(), `"'`) } if lit != "" { args = append(args, lit) } } actualCalls = append(actualCalls, args) } return true }) if len(actualCalls) != len(tt.expected) { t.Errorf("Expected %d CallExpr, got %d for command %q", len(tt.expected), len(actualCalls), tt.command) return } // Debug output to understand parser behavior t.Logf("Command: %q", tt.command) for i, call := range actualCalls { t.Logf(" CallExpr %d: %v", i, call) if len(call) > 0 { t.Logf(" args[0] = %q", call[0]) if strings.Contains(call[0], "kubectl") { t.Logf(" -> Contains kubectl!") } } } for i, expectedArgs := range tt.expected { if len(actualCalls[i]) != len(expectedArgs) { t.Errorf("CallExpr %d: expected %d args, got %d for command %q", i, len(expectedArgs), len(actualCalls[i]), tt.command) continue } for j, expectedArg := range expectedArgs { if actualCalls[i][j] != expectedArg { t.Errorf("CallExpr %d arg %d: expected %q, got %q for command %q", i, j, expectedArg, actualCalls[i][j], tt.command) } } } }) } } // TestSimplifiedKubectlDetection tests a simplified approach to kubectl detection func TestSimplifiedKubectlDetection(t *testing.T) { // Simplified kubectl detection logic isKubectl := func(args []string) bool { if len(args) == 0 { return false } // Get the first argument (the command) cmd := args[0] // Reject quoted arguments if (strings.HasPrefix(cmd, "'") && strings.HasSuffix(cmd, "'")) || (strings.HasPrefix(cmd, "\"") && strings.HasSuffix(cmd, "\"")) { return false } // Check if this is kubectl using OS-agnostic path handling basename := filepath.Base(cmd) return basename == "kubectl" || basename == "kubectl.exe" } tests := []struct { name string args []string expected bool }{ {"empty args", []string{}, false}, {"kubectl", []string{"kubectl", "get", "pods"}, true}, {"kubectl.exe", []string{"kubectl.exe", "get", "pods"}, true}, {"path to kubectl", []string{"/usr/bin/kubectl", "get", "pods"}, true}, {"Windows path", []string{"C:/tools/kubectl.exe", "delete", "pod"}, true}, {"quoted kubectl", []string{"'kubectl'", "get", "pods"}, false}, {"quoted path", []string{"\"/usr/bin/kubectl\"", "get", "pods"}, false}, {"not kubectl", []string{"ls", "-la"}, false}, {"kubectl with suffix", []string{"kubectl-1.28", "get", "pods"}, false}, {"k alias", []string{"k", "get", "pods"}, false}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := isKubectl(tt.args) if result != tt.expected { t.Errorf("isKubectl(%v) = %v, want %v", tt.args, result, tt.expected) } }) } } // TestKubectlAlwaysAtPosition0 confirms that kubectl is always at args[0] in a CallExpr func TestKubectlAlwaysAtPosition0(t *testing.T) { // Test different kubectl commands to confirm kubectl is always at position 0 testCommands := []string{ "kubectl get pods", "/usr/bin/kubectl get pods", "kubectl.exe get pods", "kubectl --context=prod get pods", "KUBECONFIG=/path/config kubectl get pods", } parser := syntax.NewParser() for _, cmd := range testCommands { t.Run(cmd, func(t *testing.T) { file, err := parser.Parse(strings.NewReader(cmd), "") if err != nil { t.Fatalf("Parse error: %v", err) } syntax.Walk(file, func(node syntax.Node) bool { if call, ok := node.(*syntax.CallExpr); ok { if len(call.Args) == 0 { return true } // Extract first argument firstArg := call.Args[0].Lit() if firstArg == "" { var sb strings.Builder syntax.NewPrinter().Print(&sb, call.Args[0]) firstArg = strings.Trim(sb.String(), `"'`) } // Check if first argument is kubectl (using same logic as main code) basename := filepath.Base(firstArg) isKubectl := basename == "kubectl" || basename == "kubectl.exe" if !isKubectl { t.Errorf("Expected kubectl at args[0], got %q (basename: %q)", firstArg, basename) } } return true }) }) } } ================================================ FILE: pkg/tools/kubectl_tool.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "fmt" "os" "strings" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" ) type Kubectl struct { executor sandbox.Executor } func NewKubectlTool(executor sandbox.Executor) *Kubectl { return &Kubectl{executor: executor} } func (t *Kubectl) Name() string { return "kubectl" } func (t *Kubectl) Description() string { return `Executes a kubectl command against the user's Kubernetes cluster. Use this tool only when you need to query or modify the state of the user's Kubernetes cluster. IMPORTANT: Interactive commands are not supported in this environment. This includes: - kubectl exec with -it flag (use non-interactive exec instead) - kubectl edit (use kubectl get -o yaml, kubectl patch, or kubectl apply instead) - kubectl port-forward (use alternative methods like NodePort or LoadBalancer) For interactive operations, please use these non-interactive alternatives: - Instead of 'kubectl edit', use 'kubectl get -o yaml' to view, 'kubectl patch' for targeted changes, or 'kubectl apply' to apply full changes - Instead of 'kubectl exec -it', use 'kubectl exec' with a specific command - Instead of 'kubectl port-forward', use service types like NodePort or LoadBalancer` } func (t *Kubectl) FunctionDefinition() *gollm.FunctionDefinition { return &gollm.FunctionDefinition{ Name: t.Name(), Description: t.Description(), Parameters: &gollm.Schema{ Type: gollm.TypeObject, Properties: map[string]*gollm.Schema{ "command": { Type: gollm.TypeString, Description: `The complete kubectl command to execute. Prefer to use heredoc syntax for multi-line commands. Please include the kubectl prefix as well. IMPORTANT: Do not use interactive commands. Instead: - Use 'kubectl get -o yaml', 'kubectl patch', or 'kubectl apply' instead of 'kubectl edit' - Use 'kubectl exec' with specific commands instead of 'kubectl exec -it' - Use service types like NodePort or LoadBalancer instead of 'kubectl port-forward' Examples: user: what pods are running in the cluster? assistant: kubectl get pods user: what is the status of the pod my-pod? assistant: kubectl get pod my-pod -o jsonpath='{.status.phase}' user: I need to edit the pod configuration assistant: # Option 1: Using patch for targeted changes kubectl patch pod my-pod --patch '{"spec":{"containers":[{"name":"main","image":"new-image"}]}}' # Option 2: Using get and apply for full changes kubectl get pod my-pod -o yaml > pod.yaml # Edit pod.yaml locally kubectl apply -f pod.yaml user: I need to execute a command in the pod assistant: kubectl exec my-pod -- /bin/sh -c "your command here"`, }, "modifies_resource": { Type: gollm.TypeString, Description: `Whether the command modifies a kubernetes resource. Possible values: - "yes" if the command modifies a resource - "no" if the command does not modify a resource - "unknown" if the command's effect on the resource is unknown`}, }, }, } } func (t *Kubectl) Run(ctx context.Context, args map[string]any) (any, error) { kubeconfig := ctx.Value(KubeconfigKey).(string) workDir := ctx.Value(WorkDirKey).(string) // Add nil check for command commandVal, ok := args["command"] if !ok || commandVal == nil { return &sandbox.ExecResult{Command: "", Error: "kubectl command not provided or is nil"}, nil } command, ok := commandVal.(string) if !ok { return &sandbox.ExecResult{Command: command, Error: "kubectl command must be a string"}, nil } // Check for interactive commands before proceeding if err := validateKubectlCommand(command); err != nil { return &sandbox.ExecResult{Command: command, Error: err.Error()}, nil } // Prepare environment env := os.Environ() if kubeconfig != "" { kubeconfig, err := ExpandShellVar(kubeconfig) if err != nil { return nil, err } env = append(env, "KUBECONFIG="+kubeconfig) } return ExecuteWithStreamingHandling(ctx, t.executor, command, workDir, env, DetectKubectlStreaming) } // DetectKubectlStreaming checks if a kubectl command is a streaming command func DetectKubectlStreaming(command string) (bool, string) { isWatch := strings.Contains(command, " get ") && strings.Contains(command, " -w") isLogs := strings.Contains(command, " logs ") && strings.Contains(command, " -f") isAttach := strings.Contains(command, " attach ") if isWatch { return true, "watch" } if isLogs { return true, "logs" } if isAttach { return true, "attach" } return false, "" } func (t *Kubectl) IsInteractive(args map[string]any) (bool, error) { commandVal, ok := args["command"] if !ok || commandVal == nil { return false, nil } command, ok := commandVal.(string) if !ok { return false, nil } return IsInteractiveCommand(command) } // CheckModifiesResource determines if the command modifies kubernetes resources // This is used for permission checks before command execution // Returns "yes", "no", or "unknown" func (t *Kubectl) CheckModifiesResource(args map[string]any) string { command, ok := args["command"].(string) if !ok { return "unknown" } return kubectlModifiesResource(command) } func validateKubectlCommand(command string) error { if strings.Contains(command, "kubectl edit") { return fmt.Errorf("interactive mode not supported for kubectl, please use non-interactive commands") } if strings.Contains(command, "kubectl port-forward") { return fmt.Errorf("port-forwarding is not allowed because assistant is running in an unattended mode, please try some other alternative") } return nil } ================================================ FILE: pkg/tools/mcp_tool.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Package tools implements the kubectl-ai tool system. package tools import ( "context" "fmt" "github.com/GoogleCloudPlatform/kubectl-ai/gollm" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/mcp" "k8s.io/klog/v2" ) // ============================================================================= // Schema Conversion Functions (kubectl-ai specific) // ============================================================================= // ConvertToolToGollm converts an MCP tool to gollm.FunctionDefinition with a simple schema func ConvertToolToGollm(mcpTool *mcp.Tool) (*gollm.FunctionDefinition, error) { def := &gollm.FunctionDefinition{ Name: mcpTool.Name, Description: mcpTool.Description, Parameters: mcpTool.InputSchema, } return def, nil } // ============================================================================= // MCP Tool Implementation // ============================================================================= // MCPTool wraps an MCP server tool to implement the Tool interface. // It serves as an adapter between MCP-based tools and kubectl-ai's tool system. type MCPTool struct { serverName string toolName string description string schema *gollm.FunctionDefinition manager *mcp.Manager } // NewMCPTool creates a new MCP tool wrapper. func NewMCPTool(serverName, toolName, description string, schema *gollm.FunctionDefinition, manager *mcp.Manager) *MCPTool { return &MCPTool{ serverName: serverName, toolName: toolName, description: description, schema: schema, manager: manager, } } // Name returns the tool name. func (t *MCPTool) Name() string { return t.toolName } func (t *MCPTool) UniqueToolName() string { return fmt.Sprintf("%s_%s", t.serverName, t.toolName) } // ServerName returns the MCP server name. func (t *MCPTool) ServerName() string { return t.serverName } // Description returns the tool description. func (t *MCPTool) Description() string { return t.description } // FunctionDefinition returns the tool's function definition. func (t *MCPTool) FunctionDefinition() *gollm.FunctionDefinition { return t.schema } // TODO(tuannvm): This is a placeholder implementation. Need to implement detection of interactive MCP tools. // IsInteractive checks if the tool requires interactive input. func (t *MCPTool) IsInteractive(args map[string]any) (bool, error) { return false, nil } // CheckModifiesResource determines if the command modifies kubernetes resources // For MCP tools, we'll conservatively assume they might modify resources // since we can't easily determine this for arbitrary external tools // Returns "yes", "no", or "unknown" func (t *MCPTool) CheckModifiesResource(args map[string]any) string { // Since MCP tools can be arbitrary external tools and we don't have a way to know // if they modify resources, we'll conservatively return "unknown" return "unknown" } // Run executes the MCP tool by calling the appropriate MCP server. func (t *MCPTool) Run(ctx context.Context, args map[string]any) (any, error) { log := klog.FromContext(ctx) // Get MCP client for the server client, exists := t.manager.GetClient(t.serverName) if !exists { return nil, fmt.Errorf("MCP server %q not connected", t.serverName) } // // Convert arguments to proper types for MCP server using the MCP package's functions // args = mcp.ConvertArgs(args) // Execute tool on MCP server result, err := client.CallTool(ctx, t.toolName, args) if err != nil { log.Info("tool info", "name", t.toolName, "schema", t.schema) log.Info("call info", "args", args) return nil, fmt.Errorf("calling MCP tool %q on server %q: %w", t.toolName, t.serverName, err) } return result, nil } ================================================ FILE: pkg/tools/streaming.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" ) // StreamDetector determines if a command is a streaming command and returns the stream type. // It returns (true, streamType) if it is a streaming command, and (false, "") otherwise. type StreamDetector func(command string) (isStreaming bool, streamType string) // ExecuteWithStreamingHandling executes a command using the provided executor, // handling streaming commands (watch, logs -f, attach) by applying a timeout // and capturing partial output. func ExecuteWithStreamingHandling(ctx context.Context, executor sandbox.Executor, command string, workDir string, env []string, detector StreamDetector) (*sandbox.ExecResult, error) { isStreaming, streamType := false, "" if detector != nil { isStreaming, streamType = detector(command) } var cmdCtx context.Context var cancel context.CancelFunc if isStreaming { // Create a context with timeout for streaming commands cmdCtx, cancel = context.WithTimeout(ctx, 7*time.Second) defer cancel() } else { // Use the provided context directly cmdCtx = ctx cancel = func() {} // No-op cancel } result, err := executor.Execute(cmdCtx, command, env, workDir) // If executor returns nil result on error (it shouldn't, but let's be safe), create one if result == nil { result = &sandbox.ExecResult{Command: command} } if isStreaming { if cmdCtx.Err() == context.DeadlineExceeded { // Timeout is expected for streaming commands result.StreamType = "timeout" result.Error = "Timeout reached after 7 seconds" // Clear the error if it was just the timeout err = nil // Set the detected stream type result.StreamType = streamType return result, nil } } return result, err } ================================================ FILE: pkg/tools/tools.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package tools import ( "context" "encoding/json" "fmt" "maps" "os" "path/filepath" "runtime" "slices" "sort" "strings" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox" "github.com/google/uuid" "sigs.k8s.io/yaml" ) type ContextKey string const ( KubeconfigKey ContextKey = "kubeconfig" WorkDirKey ContextKey = "work_dir" ExecutorKey ContextKey = "executor" ) func Lookup(name string) Tool { return allTools.Lookup(name) } var allTools Tools = Tools{ tools: make(map[string]Tool), } func Default() Tools { return allTools } // RegisterTool makes a tool available to the LLM. func RegisterTool(tool Tool) { allTools.RegisterTool(tool) } type Tools struct { tools map[string]Tool } func (t *Tools) Init() { t.tools = make(map[string]Tool) } func (t *Tools) Lookup(name string) Tool { return t.tools[name] } func (t *Tools) AllTools() []Tool { return slices.Collect(maps.Values(t.tools)) } func (t *Tools) Names() []string { names := make([]string, 0, len(t.tools)) for name := range t.tools { names = append(names, name) } sort.Strings(names) return names } func (t *Tools) RegisterTool(tool Tool) { // if mcp tool use unique name name := tool.Name() // For MCP tools, we need to use a unique name to avoid conflicts // with built-in tools or tools from other MCP servers. if mcpTool, ok := tool.(*MCPTool); ok { name = mcpTool.UniqueToolName() } if _, exists := t.tools[name]; exists { panic("tool already registered: " + name) } t.tools[name] = tool } // CloneWithExecutor creates a shallow copy of the Tools collection, // but clones any tools that need a session-specific executor (like CustomTool). func (t *Tools) CloneWithExecutor(executor sandbox.Executor) Tools { newTools := Tools{ tools: make(map[string]Tool), } for name, tool := range t.tools { // If it's a CustomTool, we need to clone it with the session-specific executor if ct, ok := tool.(*CustomTool); ok { newTools.tools[name] = ct.CloneWithExecutor(executor) } else { // For other tools (like MCP tools), we reuse the existing instance newTools.tools[name] = tool } } return newTools } type ToolCall struct { tool Tool name string arguments map[string]any } // Description returns a description of the tool call. // This is used to display the tool call in the UI. // It should be human-readable, // and should be concise enough that the user can read it quickly, // but precise enough that the user can decide whether to invoke the tool. func (t *ToolCall) Description() string { // Check if this is an MCP tool and format accordingly if mcpTool, ok := t.tool.(*MCPTool); ok { if command, ok := t.arguments["command"]; ok { return fmt.Sprintf("[MCP: %s] %s", mcpTool.serverName, command.(string)) } var args []string for k, v := range t.arguments { args = append(args, fmt.Sprintf("%s=%v", k, v)) } sort.Strings(args) return fmt.Sprintf("[MCP: %s] %s(%s)", mcpTool.serverName, t.name, strings.Join(args, ", ")) } // Default formatting for non-MCP tools if command, ok := t.arguments["command"]; ok { return command.(string) } var args []string for k, v := range t.arguments { args = append(args, fmt.Sprintf("%s=%v", k, v)) } sort.Strings(args) return fmt.Sprintf("%s(%s)", t.name, strings.Join(args, ", ")) } // ParseToolInvocation parses a request from the LLM into a tool call. func (t *Tools) ParseToolInvocation(ctx context.Context, name string, arguments map[string]any) (*ToolCall, error) { tool := t.Lookup(name) if tool == nil { return nil, fmt.Errorf("tool %q not recognized", name) } return &ToolCall{ tool: tool, name: name, arguments: arguments, }, nil } type InvokeToolOptions struct { WorkDir string // Kubeconfig is the path to the kubeconfig file. Kubeconfig string // Executor is the executor for tool execution Executor sandbox.Executor } type ToolRequestEvent struct { CallID string `json:"id,omitempty"` Name string `json:"name,omitempty"` Arguments map[string]any `json:"arguments,omitempty"` } type ToolResponseEvent struct { CallID string `json:"id,omitempty"` Response any `json:"response,omitempty"` Error string `json:"error,omitempty"` } // InvokeTool handles the execution of a single action func (t *ToolCall) InvokeTool(ctx context.Context, opt InvokeToolOptions) (any, error) { recorder := journal.RecorderFromContext(ctx) callID := uuid.NewString() recorder.Write(ctx, &journal.Event{ Timestamp: time.Now(), Action: "tool-request", Payload: ToolRequestEvent{ CallID: callID, Name: t.name, Arguments: t.arguments, }, }) ctx = context.WithValue(ctx, KubeconfigKey, opt.Kubeconfig) ctx = context.WithValue(ctx, WorkDirKey, opt.WorkDir) if opt.Executor != nil { ctx = context.WithValue(ctx, ExecutorKey, opt.Executor) } response, err := t.tool.Run(ctx, t.arguments) { ev := ToolResponseEvent{ CallID: callID, Response: response, } if err != nil { ev.Error = err.Error() } recorder.Write(ctx, &journal.Event{ Timestamp: time.Now(), Action: "tool-response", Payload: ev, }) } return response, err } // ToolResultToMap converts an arbitrary result to a map[string]any func ToolResultToMap(result any) (map[string]any, error) { // Handle simple string results (common with MCP tools) if str, ok := result.(string); ok { return map[string]any{"content": str}, nil } // Handle nil results if result == nil { return map[string]any{"content": ""}, nil } // Try to convert to map via JSON for structured results b, err := json.Marshal(result) if err != nil { return nil, fmt.Errorf("converting result to json: %w", err) } m := make(map[string]any) if err := json.Unmarshal(b, &m); err != nil { // If JSON unmarshal fails, wrap the original result return map[string]any{"content": result}, nil } return m, nil } // LoadAndRegisterCustomTools loads tool configurations from a YAML file // and registers them. func LoadAndRegisterCustomTools(configPath string) error { pathInfo, err := os.Stat(configPath) if err != nil { return fmt.Errorf("failed to describe config file %s: %w", configPath, err) } if pathInfo.IsDir() { configPaths, err := os.ReadDir(configPath) if err != nil { return fmt.Errorf("failed to read config dir %s: %w", configPath, err) } for _, entry := range configPaths { if err := LoadAndRegisterCustomTools(filepath.Join(configPath, entry.Name())); err != nil { return err } } return nil } yamlFile, err := os.ReadFile(configPath) if os.IsNotExist(err) { return nil } else if err != nil { return fmt.Errorf("failed to read config file %s: %w", configPath, err) } var configs []CustomToolConfig err = yaml.Unmarshal(yamlFile, &configs) if err != nil { return fmt.Errorf("failed to parse YAML config file %s: %w", configPath, err) } // Register each custom tool var registrationErrors []string for _, config := range configs { tool, err := NewCustomTool(config) if err != nil { registrationErrors = append(registrationErrors, fmt.Sprintf("failed to create tool %q: %v", config.Name, err)) continue // Skip registration if creation failed } // Check for duplicate registration attempt if _, exists := allTools.tools[tool.Name()]; exists { registrationErrors = append(registrationErrors, fmt.Sprintf("tool %q already registered (possibly built-in), skipping custom definition", tool.Name())) continue } RegisterTool(tool) } if len(registrationErrors) > 0 { return fmt.Errorf("encountered errors during custom tool registration:\n - %s", strings.Join(registrationErrors, "\n - ")) } return nil } // For CustomTool func (t *CustomTool) IsInteractive(args map[string]any) (bool, error) { // Custom tools are not interactive by default return false, nil } // Add a method to access the tool func (t *ToolCall) GetTool() Tool { return t.tool } // ExpandShellVar expands shell variables and syntax using bash func ExpandShellVar(value string) (string, error) { if strings.Contains(value, "~") { if len(value) >= 2 && value[0] == '~' && os.IsPathSeparator(value[1]) { if runtime.GOOS == "windows" { value = filepath.Join(os.Getenv("USERPROFILE"), value[2:]) } else { value = filepath.Join(os.Getenv("HOME"), value[2:]) } } } return os.ExpandEnv(value), nil } func IsInteractiveCommand(command string) (bool, error) { // Inline isKubectlCommand logic words := strings.Fields(command) if len(words) == 0 { return false, nil } base := filepath.Base(words[0]) if base != "kubectl" { return false, nil } isExec := strings.Contains(command, " exec ") && strings.Contains(command, " -it") isPortForward := strings.Contains(command, " port-forward ") isEdit := strings.Contains(command, " edit ") if isExec || isPortForward || isEdit { return true, fmt.Errorf("interactive mode not supported for kubectl, please use non-interactive commands") } return false, nil } ================================================ FILE: pkg/ui/html/htmlui.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package html import ( "context" _ "embed" "encoding/json" "errors" "fmt" "net" "net/http" "os" "strconv" "sync" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/agent" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/sessions" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/ui" "github.com/charmbracelet/glamour" "golang.org/x/sync/errgroup" "k8s.io/klog/v2" ) // Broadcaster manages a set of clients for Server-Sent Events. type Broadcaster struct { clients map[chan []byte]bool newClient chan chan []byte delClient chan chan []byte messages chan []byte mu sync.Mutex } // NewBroadcaster creates a new Broadcaster instance. func NewBroadcaster() *Broadcaster { b := &Broadcaster{ clients: make(map[chan []byte]bool), newClient: make(chan (chan []byte)), delClient: make(chan (chan []byte)), messages: make(chan []byte, 10), } return b } // Run starts the broadcaster's event loop. func (b *Broadcaster) Run(ctx context.Context) { for { select { case <-ctx.Done(): return case client := <-b.newClient: b.mu.Lock() b.clients[client] = true b.mu.Unlock() case client := <-b.delClient: b.mu.Lock() delete(b.clients, client) close(client) b.mu.Unlock() case msg := <-b.messages: b.mu.Lock() for client := range b.clients { select { case client <- msg: default: klog.Warning("SSE client buffer full, dropping message.") } } b.mu.Unlock() } } } // Broadcast sends a message to all connected clients. func (b *Broadcaster) Broadcast(msg []byte) { b.messages <- msg } type HTMLUserInterface struct { httpServer *http.Server httpServerListener net.Listener manager *agent.AgentManager sessionManager *sessions.SessionManager journal journal.Recorder defaultModel string defaultProvider string markdownRenderer *glamour.TermRenderer broadcasters map[string]*Broadcaster broadcastersMu sync.Mutex broadcasterCancels map[string]context.CancelFunc baseCtx context.Context } var _ ui.UI = &HTMLUserInterface{} func NewHTMLUserInterface(manager *agent.AgentManager, sessionManager *sessions.SessionManager, defaultModel, defaultProvider string, listenAddress string, journal journal.Recorder) (*HTMLUserInterface, error) { mux := http.NewServeMux() u := &HTMLUserInterface{ manager: manager, sessionManager: sessionManager, defaultModel: defaultModel, defaultProvider: defaultProvider, journal: journal, broadcasters: make(map[string]*Broadcaster), broadcasterCancels: make(map[string]context.CancelFunc), } // Register callback to listen to new agents manager.SetAgentCreatedCallback(func(a *agent.Agent) { u.ensureAgentListener(a) }) httpServer := &http.Server{ Addr: listenAddress, Handler: mux, } mux.HandleFunc("GET /", u.serveIndex) mux.HandleFunc("GET /api/sessions", u.handleListSessions) mux.HandleFunc("POST /api/sessions", u.handleCreateSession) mux.HandleFunc("POST /api/sessions/{id}/rename", u.handleRenameSession) mux.HandleFunc("DELETE /api/sessions/{id}", u.handleDeleteSession) mux.HandleFunc("GET /api/sessions/{id}/stream", u.handleSessionStream) mux.HandleFunc("POST /api/sessions/{id}/send-message", u.handlePOSTSendMessage) mux.HandleFunc("POST /api/sessions/{id}/choose-option", u.handlePOSTChooseOption) httpServerListener, err := net.Listen("tcp", listenAddress) if err != nil { return nil, fmt.Errorf("starting http server network listener: %w", err) } endpoint := httpServerListener.Addr() u.httpServerListener = httpServerListener u.httpServer = httpServer fmt.Fprintf(os.Stdout, "listening on http://%s\n", endpoint) mdRenderer, err := glamour.NewTermRenderer( glamour.WithAutoStyle(), glamour.WithPreservedNewLines(), glamour.WithEmoji(), ) if err != nil { return nil, fmt.Errorf("error initializing the markdown renderer: %w", err) } u.markdownRenderer = mdRenderer return u, nil } func (u *HTMLUserInterface) Run(ctx context.Context) error { g, gctx := errgroup.WithContext(ctx) u.baseCtx = gctx g.Go(func() error { if err := u.httpServer.Serve(u.httpServerListener); err != nil && !errors.Is(err, http.ErrServerClosed) { return fmt.Errorf("error running http server: %w", err) } return nil }) g.Go(func() error { <-gctx.Done() shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := u.httpServer.Shutdown(shutdownCtx); err != nil { klog.Errorf("HTTP server shutdown error: %v", err) } return nil }) return g.Wait() } //go:embed index.html var indexHTML []byte func (u *HTMLUserInterface) serveIndex(w http.ResponseWriter, req *http.Request) { w.Header().Set("Content-Type", "text/html") w.Write(indexHTML) } func (u *HTMLUserInterface) handleSessionStream(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) id := req.PathValue("id") if id == "" { http.Error(w, "missing session id", http.StatusBadRequest) return } flusher, ok := w.(http.Flusher) if !ok { http.Error(w, "Streaming unsupported!", http.StatusInternalServerError) return } w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") clientChan := make(chan []byte, 10) broadcaster := u.getBroadcaster(id) broadcaster.newClient <- clientChan defer func() { broadcaster.delClient <- clientChan }() log.Info("SSE client connected", "sessionID", id) agent, err := u.manager.GetAgent(ctx, id) var initialData []byte if err != nil { log.Error(err, "getting agent for session") } else { initialData, err = u.getSessionStateJSON(agent.Session) } if err != nil { log.Error(err, "getting initial state for SSE client") } else { fmt.Fprintf(w, "data: %s\n\n", initialData) flusher.Flush() } for { select { case <-ctx.Done(): log.Info("SSE client disconnected") return case msg := <-clientChan: fmt.Fprintf(w, "data: %s\n\n", msg) flusher.Flush() } } } func (u *HTMLUserInterface) handleListSessions(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) sessionsList, err := u.manager.ListSessions() if err != nil { log.Error(err, "listing sessions") http.Error(w, err.Error(), http.StatusInternalServerError) return } w.Header().Set("Content-Type", "application/json") if err := json.NewEncoder(w).Encode(sessionsList); err != nil { log.Error(err, "encoding sessions list") } } func (u *HTMLUserInterface) handleCreateSession(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) meta := sessions.Metadata{ ModelID: u.defaultModel, ProviderID: u.defaultProvider, } session, err := u.sessionManager.NewSession(meta) if err != nil { log.Error(err, "creating new session") http.Error(w, err.Error(), http.StatusInternalServerError) return } // Ensure agent is started/loaded (though mostly for side effect of starting if not started) if _, err := u.manager.GetAgent(ctx, session.ID); err != nil { log.Error(err, "starting agent for new session") // We don't fail the request here necessarily, but it's good to know. } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]string{"id": session.ID}) } func (u *HTMLUserInterface) handleRenameSession(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) id := req.PathValue("id") if id == "" { http.Error(w, "missing session id", http.StatusBadRequest) return } if err := req.ParseForm(); err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } newName := req.FormValue("name") if newName == "" { http.Error(w, "missing name", http.StatusBadRequest) return } session, err := u.manager.FindSessionByID(id) if err != nil { http.Error(w, "session not found", http.StatusNotFound) return } session.Name = newName if err := u.manager.UpdateLastAccessed(session); err != nil { // UpdateLastAccessed also saves the session log.Error(err, "updating session") http.Error(w, err.Error(), http.StatusInternalServerError) return } if agent, err := u.manager.GetAgent(ctx, id); err == nil { agent.Session.Name = newName // Broadcast update if data, err := u.getSessionStateJSON(agent.Session); err == nil { u.getBroadcaster(id).Broadcast(data) } } w.WriteHeader(http.StatusOK) } func (u *HTMLUserInterface) handleDeleteSession(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) id := req.PathValue("id") if id == "" { http.Error(w, "missing session id", http.StatusBadRequest) return } if err := u.manager.DeleteSession(id); err != nil { log.Error(err, "deleting session") http.Error(w, err.Error(), http.StatusInternalServerError) return } // If anyone was listening to this session, they should know it's gone. // We can close the broadcaster. u.broadcastersMu.Lock() if cancel, ok := u.broadcasterCancels[id]; ok { cancel() delete(u.broadcasterCancels, id) } delete(u.broadcasters, id) u.broadcastersMu.Unlock() w.WriteHeader(http.StatusOK) } func (u *HTMLUserInterface) handlePOSTSendMessage(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) id := req.PathValue("id") if id == "" { http.Error(w, "missing session id", http.StatusBadRequest) return } if err := req.ParseForm(); err != nil { log.Error(err, "parsing form") http.Error(w, err.Error(), http.StatusBadRequest) return } q := req.FormValue("q") if q == "" { http.Error(w, "missing query", http.StatusBadRequest) return } // Get the agent for this session agent, err := u.manager.GetAgent(ctx, id) if err != nil { log.Error(err, "getting agent") http.Error(w, err.Error(), http.StatusInternalServerError) return } // Send the message to the agent agent.Input <- &api.UserInputResponse{Query: q} w.WriteHeader(http.StatusOK) } func (u *HTMLUserInterface) handlePOSTChooseOption(w http.ResponseWriter, req *http.Request) { ctx := req.Context() log := klog.FromContext(ctx) id := req.PathValue("id") if id == "" { http.Error(w, "missing session id", http.StatusBadRequest) return } if err := req.ParseForm(); err != nil { log.Error(err, "parsing form") http.Error(w, err.Error(), http.StatusBadRequest) return } choice := req.FormValue("choice") if choice == "" { http.Error(w, "missing choice", http.StatusBadRequest) return } choiceIndex, err := strconv.Atoi(choice) if err != nil { http.Error(w, "invalid choice", http.StatusBadRequest) return } // Get the agent agent, err := u.manager.GetAgent(ctx, id) if err != nil { http.Error(w, "agent not found", http.StatusNotFound) return } // Send the choice to the agent agent.Input <- &api.UserChoiceResponse{Choice: choiceIndex} w.WriteHeader(http.StatusOK) } func (u *HTMLUserInterface) Close() error { var errs []error if u.httpServerListener != nil { if err := u.httpServerListener.Close(); err != nil { errs = append(errs, err) } else { u.httpServerListener = nil } } u.broadcastersMu.Lock() for id, cancel := range u.broadcasterCancels { cancel() delete(u.broadcasterCancels, id) } u.broadcasters = make(map[string]*Broadcaster) u.broadcastersMu.Unlock() return errors.Join(errs...) } func (u *HTMLUserInterface) ClearScreen() { // Not applicable for HTML UI } func (u *HTMLUserInterface) getSessionStateJSON(session *api.Session) ([]byte, error) { allMessages := session.AllMessages() // Create a copy of the messages to avoid race conditions var messages []*api.Message for _, message := range allMessages { if message.Type == api.MessageTypeUserInputRequest && message.Payload == ">>>" { continue } messages = append(messages, message) } agentState := session.AgentState data := map[string]interface{}{ "messages": messages, "agentState": agentState, "sessionId": session.ID, } return json.Marshal(data) } func (u *HTMLUserInterface) getBroadcaster(sessionID string) *Broadcaster { u.broadcastersMu.Lock() defer u.broadcastersMu.Unlock() if b, ok := u.broadcasters[sessionID]; ok { return b } b := NewBroadcaster() u.broadcasters[sessionID] = b parent := u.baseCtx if parent == nil { parent = context.Background() } ctx, cancel := context.WithCancel(parent) u.broadcasterCancels[sessionID] = cancel // Start the broadcaster loop go b.Run(ctx) return b } func (u *HTMLUserInterface) ensureAgentListener(a *agent.Agent) { // Start a goroutine to listen to this agent's output go func() { for range a.Output { // Broadcast state if a.Session == nil { continue } data, err := u.getSessionStateJSON(a.Session) if err != nil { klog.Errorf("Error marshaling state for broadcast: %v", err) continue } b := u.getBroadcaster(a.Session.ID) b.Broadcast(data) } }() } ================================================ FILE: pkg/ui/html/index.html ================================================ kubectl-ai
================================================ FILE: pkg/ui/interfaces.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package ui import ( "context" "fmt" ) // UI is the interface that defines the capabilities of assisant's user interface. // Each of the UIs, CLI, TUI, Web, etc. implement this interface. type UI interface { // ClearScreen clears any output rendered to the screen ClearScreen() // Run starts the UI and blocks until the context is done. Run(ctx context.Context) error } // Type is the type of user interface. type Type string const ( UITypeTerminal Type = "terminal" UITypeWeb Type = "web" UITypeTUI Type = "tui" ) // Implement pflag.Value for UIType func (u *Type) Set(s string) error { switch s { case "terminal", "web", "tui": *u = Type(s) return nil default: return fmt.Errorf("invalid UI type: %s", s) } } func (u *Type) String() string { return string(*u) } func (u *Type) Type() string { return "UIType" } ================================================ FILE: pkg/ui/terminal.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package ui import ( "bufio" "context" "encoding/json" "errors" "fmt" "io" "os" "path/filepath" "strconv" "strings" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/agent" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools" "github.com/charmbracelet/glamour" "github.com/chzyer/readline" "golang.org/x/term" "k8s.io/klog/v2" ) type computedStyle struct { Foreground colorValue RenderMarkdown bool } type colorValue string const ( colorGreen colorValue = "green" colorWhite colorValue = "white" colorRed colorValue = "red" ) type styleOption func(s *computedStyle) func foreground(color colorValue) styleOption { return func(s *computedStyle) { s.Foreground = color } } func renderMarkdown() styleOption { return func(s *computedStyle) { s.RenderMarkdown = true } } // TODO: rename this to CLI because the command line interface. type TerminalUI struct { journal journal.Recorder markdownRenderer *glamour.TermRenderer // Input handling fields (initialized once) rlInstance *readline.Instance // For readline input ttyFile *os.File // For TTY input ttyReaderInstance *bufio.Reader // For TTY input // This is useful in cases where stdin is already been used for providing the input to the agent (caller in this case) // in such cases, stdin is already consumed and closed and reading input results in IO error. // In such cases, we open /dev/tty and use it for taking input. useTTYForInput bool // showToolOutput disables truncation of tool output. showToolOutput bool agent *agent.Agent } var _ UI = &TerminalUI{} func getCustomTerminalWidth() int { // Check for user-configured width via environment variable if widthStr := os.Getenv("KUBECTL_AI_TERM_WIDTH"); widthStr != "" { if widthStr == "auto" { width, _, err := term.GetSize(int(os.Stdout.Fd())) if err != nil { klog.Warningf("Failed to get terminal size: %v, using default width", err) return 0 } return width } if width, err := strconv.Atoi(widthStr); err == nil && width > 0 { return width } klog.Warningf("Invalid KUBECTL_AI_TERM_WIDTH value %q, using default", widthStr) } // Return 0 to indicate no custom width should be set (use glamour's default) return 0 } func NewTerminalUI(agent *agent.Agent, useTTYForInput bool, showToolOutput bool, journal journal.Recorder) (*TerminalUI, error) { width := getCustomTerminalWidth() options := []glamour.TermRendererOption{ glamour.WithAutoStyle(), glamour.WithPreservedNewLines(), glamour.WithEmoji(), } // Only add WordWrap if a valid width is configured if width > 0 { options = append(options, glamour.WithWordWrap(width)) } mdRenderer, err := glamour.NewTermRenderer(options...) if err != nil { return nil, fmt.Errorf("error initializing the markdown renderer: %w", err) } u := &TerminalUI{ markdownRenderer: mdRenderer, journal: journal, useTTYForInput: useTTYForInput, // Store this flag agent: agent, showToolOutput: showToolOutput, } return u, nil } func (u *TerminalUI) Run(ctx context.Context) error { session := u.agent.GetSession() if len(session.Messages) > 0 { greeting := "Welcome back. What can I help you with today?\n (Don't want to continue your last session? Use --new-session)" // If it's a persistent session (not memory), print metadata if u.agent.SessionBackend == "filesystem" { greeting = fmt.Sprintf("%s\n\n%s", greeting, session.String()) } out, _ := u.markdownRenderer.Render(greeting) fmt.Printf("\n%s\n", out) } // Channel to signal when the agent has exited agentExited := make(chan struct{}) // Start a goroutine to handle agent output go func() { for { select { case <-ctx.Done(): return case msg, ok := <-u.agent.Output: if !ok { return } klog.Infof("agent output: %+v", msg) u.handleMessage(msg.(*api.Message)) // Check if agent has exited in RunOnce mode if u.agent.GetSession().AgentState == api.AgentStateExited { klog.Info("Agent has exited, terminating UI") close(agentExited) return } } } }() // Block until context is cancelled or agent exits select { case <-ctx.Done(): return nil case <-agentExited: return u.agent.LastErr() } } func (u *TerminalUI) ttyReader() (*bufio.Reader, error) { if u.ttyReaderInstance != nil { return u.ttyReaderInstance, nil } // Initialize TTY input tty, err := os.OpenFile("/dev/tty", os.O_RDWR, 0) if err != nil { return nil, fmt.Errorf("opening tty for input: %w", err) } u.ttyFile = tty // Store file handle for closing u.ttyReaderInstance = bufio.NewReader(tty) return u.ttyReaderInstance, nil } func (u *TerminalUI) readlineInstance() (*readline.Instance, error) { if u.rlInstance != nil { return u.rlInstance, nil } // Initialize readline input historyPath := filepath.Join(os.TempDir(), "kubectl-ai-history") rl, err := readline.NewEx(&readline.Config{ Prompt: ">>> ", // Default prompt for main input Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr, HistoryFile: historyPath, // History enabled by default }) if err != nil { // Log warning or fallback if readline init fails? klog.Warningf("Failed to initialize readline, input might be limited: %v", err) // Proceed without readline for now, or return error? // Returning error to make it explicit return nil, fmt.Errorf("creating readline instance: %w", err) } u.rlInstance = rl // Store readline instance return u.rlInstance, nil } func (u *TerminalUI) Close() error { var errs []error // Close the initialized input handler if u.rlInstance != nil { if err := u.rlInstance.Close(); err != nil { errs = append(errs, fmt.Errorf("closing readline instance: %w", err)) } } if u.ttyFile != nil { if err := u.ttyFile.Close(); err != nil { errs = append(errs, fmt.Errorf("closing tty file: %w", err)) } } return errors.Join(errs...) } func (u *TerminalUI) handleMessage(msg *api.Message) { text := "" var styleOptions []styleOption switch msg.Type { case api.MessageTypeText: text = msg.Payload.(string) switch msg.Source { case api.MessageSourceUser: // styleOptions = append(styleOptions, Foreground(ColorWhite)) // since we print the message as user types, we don't need to print it again return case api.MessageSourceAgent: styleOptions = append(styleOptions, renderMarkdown(), foreground(colorGreen)) case api.MessageSourceModel: styleOptions = append(styleOptions, renderMarkdown()) } case api.MessageTypeError: styleOptions = append(styleOptions, foreground(colorRed)) text = msg.Payload.(string) case api.MessageTypeToolCallRequest: styleOptions = append(styleOptions, foreground(colorGreen)) text = fmt.Sprintf("\n Running: %s\n", msg.Payload.(string)) case api.MessageTypeToolCallResponse: if !u.showToolOutput { return } styleOptions = append(styleOptions, renderMarkdown()) output, err := tools.ToolResultToMap(msg.Payload) if err != nil { klog.Errorf("Error converting tool result to map: %v", err) u.agent.Input <- fmt.Errorf("error converting tool result to map: %w", err) return } responseText := formatToolCallResponse(output) text = fmt.Sprintf("%s\n", responseText) case api.MessageTypeUserInputRequest: text = msg.Payload.(string) klog.Infof("Received user input request with payload: %q", text) var query string if u.useTTYForInput { tReader, err := u.ttyReader() if err != nil { klog.Errorf("Failed to get TTY reader: %v", err) return } // keep reading input until we get a non-empty query for { var err error fmt.Print("\n>>> ") // Print prompt manually query, err = tReader.ReadString('\n') if err != nil { klog.Infof("TTY read error: %v", err) if err == io.EOF { // Handle Ctrl+D gracefully u.agent.Input <- io.EOF return } klog.Errorf("Error reading from TTY: %v", err) u.agent.Input <- fmt.Errorf("error reading from TTY: %w", err) return } if strings.TrimSpace(query) == "" { continue } break } klog.Infof("Sending TTY input to agent: %q", query) u.agent.Input <- &api.UserInputResponse{Query: query} } else { rlInstance, err := u.readlineInstance() if err != nil { klog.Errorf("Failed to create readline instance: %v", err) u.agent.Input <- fmt.Errorf("error creating readline instance: %w", err) return } // keep reading input until we get a non-empty query for { rlInstance.SetPrompt(">>> ") // Ensure correct prompt query, err = rlInstance.Readline() if err != nil { klog.Infof("Readline error: %v", err) switch err { case readline.ErrInterrupt: // Handle Ctrl+C u.agent.Input <- io.EOF case io.EOF: // Handle Ctrl+D u.agent.Input <- io.EOF default: u.agent.Input <- err } break } if strings.TrimSpace(query) == "" { continue } klog.Infof("Sending readline input to agent: %q", query) u.agent.Input <- &api.UserInputResponse{Query: query} break } } if query == "clear" || query == "reset" { u.ClearScreen() } return case api.MessageTypeUserChoiceRequest: choiceRequest := msg.Payload.(*api.UserChoiceRequest) prompt, _ := u.markdownRenderer.Render(choiceRequest.Prompt) fmt.Printf("\n%s\n", string(prompt)) for i, option := range choiceRequest.Options { fmt.Printf(" %d. %s\n", i+1, option.Label) } fmt.Println() var choice int for { var line string var err error if u.useTTYForInput { tReader, err := u.ttyReader() if err != nil { klog.Errorf("Failed to get TTY reader: %v", err) return } fmt.Print("Enter your choice: ") line, err = tReader.ReadString('\n') if err != nil { klog.Infof("TTY read error: %v", err) if err == io.EOF { // Handle Ctrl+D gracefully u.agent.Input <- io.EOF return } klog.Errorf("Error reading from TTY: %v", err) u.agent.Input <- fmt.Errorf("error reading from TTY: %w", err) return } } else { rlInstance, err := u.readlineInstance() if err != nil { klog.Errorf("Failed to create readline instance: %v", err) u.agent.Input <- fmt.Errorf("error creating readline instance: %w", err) return } rlInstance.SetPrompt("Enter your choice: ") line, err = rlInstance.Readline() if err != nil { klog.Infof("Readline error: %v", err) switch err { case readline.ErrInterrupt, io.EOF: u.agent.Input <- io.EOF return default: u.agent.Input <- err return } } } input := strings.TrimSpace(strings.ToLower(line)) choice = -1 // Handle special cases for yes/no if input == "y" || input == "yes" { input = "1" } if input == "n" || input == "no" { input = "3" } choiceIdx, err := strconv.Atoi(input) if err == nil && choiceIdx > 0 && choiceIdx <= len(choiceRequest.Options) { choice = choiceIdx break } fmt.Println("Invalid choice. Please try again.") } u.agent.Input <- &api.UserChoiceResponse{Choice: choice} return default: klog.Warningf("unsupported message type: %v", msg.Type) return } computedStyle := &computedStyle{} for _, opt := range styleOptions { opt(computedStyle) } printText := text if computedStyle.RenderMarkdown && printText != "" { out, err := u.markdownRenderer.Render(printText) if err != nil { klog.Errorf("Error rendering markdown: %v", err) } else { printText = out } } reset := "" switch computedStyle.Foreground { case colorRed: fmt.Printf("\033[31m") reset += "\033[0m" case colorGreen: fmt.Printf("\033[32m") reset += "\033[0m" case colorWhite: fmt.Printf("\033[37m") reset += "\033[0m" case "": default: klog.Info("foreground color not supported by TerminalUI", "color", computedStyle.Foreground) } fmt.Printf("%s%s", printText, reset) } func (u *TerminalUI) ClearScreen() { fmt.Print("\033[H\033[2J") } func formatToolCallResponse(payload map[string]any) string { if payload == nil { return "" } if v, ok := payload["content"]; ok { return fmt.Sprint(v) } if v, ok := payload["stdout"]; ok { return fmt.Sprint(v) } if b, err := json.MarshalIndent(payload, "", " "); err == nil { return string(b) } return fmt.Sprint(payload) } ================================================ FILE: pkg/ui/tui.go ================================================ // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package ui import ( "context" "fmt" "io" "os" "strings" "sync" "time" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/agent" "github.com/GoogleCloudPlatform/kubectl-ai/pkg/api" "github.com/charmbracelet/bubbles/list" "github.com/charmbracelet/bubbles/spinner" "github.com/charmbracelet/bubbles/textinput" "github.com/charmbracelet/bubbles/viewport" tea "github.com/charmbracelet/bubbletea" "github.com/charmbracelet/glamour" "github.com/charmbracelet/lipgloss" "k8s.io/klog/v2" ) const logo = ` _ _ _ _ _ | | ___ _| |__ ___ ___| |_| | __ _(_) | |/ / | | | '_ \ / _ \/ __| __| |_____ / _ | | | <| |_| | |_) | __/ (__| |_| |_____| (_| | | |_|\_\\__,_|_.__/ \___|\___|\__|_| \__,_|_| ` // Color palette - Google Material Design colors var ( colorPrimary = lipgloss.Color("#8AB4F8") // Blue 200 colorSecondary = lipgloss.Color("#81C995") // Green 200 colorError = lipgloss.Color("#F28B82") // Red 200 colorWarning = lipgloss.Color("#FDD663") // Yellow 200 colorText = lipgloss.Color("#E8EAED") // Grey 200 colorMuted = lipgloss.Color("#9AA0A6") // Grey 500 colorDim = lipgloss.Color("#5F6368") // Grey 700 colorBgSubtle = lipgloss.Color("#303134") // Surface variant colorBgCode = lipgloss.Color("#1E1E1E") // Code background ) // Styles - consolidated for reuse var ( textStyle = lipgloss.NewStyle().Foreground(colorText) mutedStyle = lipgloss.NewStyle().Foreground(colorMuted) dimStyle = lipgloss.NewStyle().Foreground(colorDim) primaryText = lipgloss.NewStyle().Foreground(colorPrimary).Bold(true) successText = lipgloss.NewStyle().Foreground(colorSecondary).Bold(true) errorText = lipgloss.NewStyle().Foreground(colorError).Bold(true) warnText = lipgloss.NewStyle().Foreground(colorWarning).Bold(true) statusBar = lipgloss.NewStyle().Background(colorBgSubtle).Foreground(colorText) userMsg = lipgloss.NewStyle(). BorderLeft(true).BorderStyle(lipgloss.ThickBorder()). BorderForeground(colorPrimary).PaddingLeft(1).MarginBottom(1) agentMsg = lipgloss.NewStyle(). BorderLeft(true).BorderStyle(lipgloss.ThickBorder()). BorderForeground(colorSecondary).PaddingLeft(1).MarginBottom(1) toolBox = lipgloss.NewStyle(). Border(lipgloss.RoundedBorder()).BorderForeground(colorSecondary). Padding(0, 1).MarginBottom(1) errorBox = lipgloss.NewStyle(). Border(lipgloss.RoundedBorder()).BorderForeground(colorError). Padding(0, 1).MarginBottom(1) inputBox = lipgloss.NewStyle().Border(lipgloss.RoundedBorder()).BorderForeground(colorPrimary).Padding(0, 1) inputBoxDim = lipgloss.NewStyle().Border(lipgloss.RoundedBorder()).BorderForeground(colorDim).Padding(0, 1) codeStyle = lipgloss.NewStyle().Foreground(colorText).Background(colorBgCode).Padding(0, 1) ) // List item for choice selection type item string func (i item) FilterValue() string { return "" } type itemDelegate struct{} func (d itemDelegate) Height() int { return 1 } func (d itemDelegate) Spacing() int { return 0 } func (d itemDelegate) Update(_ tea.Msg, _ *list.Model) tea.Cmd { return nil } func (d itemDelegate) Render(w io.Writer, m list.Model, idx int, li list.Item) { s, ok := li.(item) if !ok { return } if idx == m.Index() { fmt.Fprint(w, primaryText.Render("> "+string(s))) } else { fmt.Fprint(w, mutedStyle.PaddingLeft(2).Render(string(s))) } } // TUI is the terminal user interface for the agent. type TUI struct { program *tea.Program agent *agent.Agent } func NewTUI(agent *agent.Agent) *TUI { return &TUI{ program: tea.NewProgram(newModel(agent), tea.WithAltScreen(), tea.WithMouseAllMotion()), agent: agent, } } func (u *TUI) Run(ctx context.Context) error { // Suppress stderr to prevent klog from breaking TUI if devNull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0); err == nil { orig := os.Stderr os.Stderr = devNull defer func() { os.Stderr = orig; devNull.Close() }() } klog.SetOutput(io.Discard) klog.LogToStderr(false) go func() { for { select { case <-ctx.Done(): return case msg, ok := <-u.agent.Output: if !ok { return } u.program.Send(msg) } } }() _, err := u.program.Run() return err } func (u *TUI) ClearScreen() {} type sessionListMsg []api.SessionInfo func (m *model) fetchSessions() tea.Msg { sessions, err := m.agent.ListSessions() if err != nil { return api.Message{ Type: api.MessageTypeError, Payload: fmt.Sprintf("Failed to list sessions: %v", err), } } return sessionListMsg(sessions) } type tickMsg time.Time // Render cache for markdown type renderCache struct { mu sync.RWMutex cache map[string]string width int renderer *glamour.TermRenderer } func newRenderCache() *renderCache { return &renderCache{cache: make(map[string]string)} } func (rc *renderCache) get(id string) (string, bool) { rc.mu.RLock() defer rc.mu.RUnlock() v, ok := rc.cache[id] return v, ok } func (rc *renderCache) set(id, content string) { rc.mu.Lock() defer rc.mu.Unlock() rc.cache[id] = content } func (rc *renderCache) getRenderer(width int) (*glamour.TermRenderer, error) { rc.mu.Lock() defer rc.mu.Unlock() if rc.width != width { rc.cache = make(map[string]string) rc.width = width rc.renderer = nil } if rc.renderer == nil { r, err := glamour.NewTermRenderer(glamour.WithStylePath("dark"), glamour.WithWordWrap(width)) if err != nil { return nil, err } rc.renderer = r } return rc.renderer, nil } // Model state type model struct { agent *agent.Agent viewport viewport.Model input textinput.Model spinner spinner.Model list list.Model cache *renderCache messages []*api.Message width int height int dirty bool quitting bool thinkStart time.Time // Choice mode tracking inChoiceMode bool choicePrompt string choiceOptionID string // Track which choice request we initialized for choiceType string // "confirm" or "session" sessionIDs []string } func newModel(agent *agent.Agent) model { ti := textinput.New() ti.Placeholder = "Ask kubectl-ai anything..." ti.Focus() ti.Prompt = "" ti.CharLimit = 4096 ti.Width = 80 ti.TextStyle = textStyle ti.PlaceholderStyle = dimStyle ti.Cursor.Style = primaryText sp := spinner.New() sp.Spinner = spinner.MiniDot sp.Style = primaryText l := list.New(nil, itemDelegate{}, 40, 5) l.SetShowStatusBar(false) l.SetFilteringEnabled(false) l.SetShowHelp(false) l.SetShowPagination(false) l.SetShowTitle(false) vp := viewport.New(80, 20) vp.MouseWheelEnabled = true return model{ agent: agent, input: ti, viewport: vp, spinner: sp, list: l, cache: newRenderCache(), dirty: true, } } func (m model) Init() tea.Cmd { return tea.Batch(textinput.Blink, m.spinner.Tick, m.tick()) } func (m model) tick() tea.Cmd { return tea.Tick(time.Second, func(t time.Time) tea.Msg { return tickMsg(t) }) } func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { case tea.WindowSizeMsg: m.width, m.height = msg.Width, msg.Height m.dirty = true m.resize() return m, nil case tea.KeyMsg: return m.handleKey(msg) case tea.MouseMsg: switch msg.Button { case tea.MouseButtonWheelUp: m.viewport.ScrollUp(3) case tea.MouseButtonWheelDown: m.viewport.ScrollDown(3) } return m, nil case *api.Message: return m.handleAgentMsg(msg) case spinner.TickMsg: var cmd tea.Cmd m.spinner, cmd = m.spinner.Update(msg) return m, cmd case tickMsg: return m, m.tick() case sessionListMsg: if len(msg) == 0 { m.messages = append(m.messages, &api.Message{ Source: api.MessageSourceAgent, Type: api.MessageTypeText, Payload: "No sessions found.", Timestamp: time.Now(), }) m.dirty = true m.refresh() m.viewport.GotoBottom() return m, nil } items := make([]list.Item, len(msg)) ids := make([]string, len(msg)) for i, s := range msg { label := fmt.Sprintf("%s (%s) • %d msgs", s.ID, s.ModelID, s.MessageCount) if s.Name != "" { label = fmt.Sprintf("%s (%s) • %s • %d msgs", s.Name, s.ModelID, s.ID, s.MessageCount) } items[i] = item(label) ids[i] = s.ID } m.list.SetItems(items) m.list.Select(0) m.inChoiceMode = true m.choicePrompt = "Select a session to resume" m.choiceOptionID = "manual-session-picker" m.choiceType = "session" m.sessionIDs = ids m.dirty = true m.refresh() m.viewport.GotoBottom() return m, nil } return m, nil } func (m *model) resize() { m.viewport.Width = m.width - 2 m.input.Width = m.width - 6 m.list.SetWidth(m.width - 4) m.updateViewportHeight() m.refresh() m.viewport.GotoBottom() } func (m *model) updateViewportHeight() { // Layout: status(1) + 2 dividers(2) + input(3) + help(1) + bottom padding(1) = 8 contentH := m.height - 8 contentH = max(contentH, 5) m.viewport.Height = contentH } func (m *model) navigateList(keyType tea.KeyType) tea.Cmd { var cmd tea.Cmd m.list, cmd = m.list.Update(tea.KeyMsg{Type: keyType}) m.dirty = true m.refresh() return cmd } func (m *model) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { switch msg.Type { case tea.KeyCtrlC, tea.KeyCtrlD: m.quitting = true return m, tea.Quit case tea.KeyEsc: m.input.Reset() return m, nil case tea.KeyEnter: return m.handleEnter() case tea.KeyUp: if m.inChoiceMode { return m, m.navigateList(tea.KeyUp) } m.viewport.ScrollUp(1) case tea.KeyDown: if m.inChoiceMode { return m, m.navigateList(tea.KeyDown) } m.viewport.ScrollDown(1) case tea.KeyPgUp: m.viewport.ScrollUp(m.viewport.Height / 2) case tea.KeyPgDown: m.viewport.ScrollDown(m.viewport.Height / 2) default: switch msg.String() { case "ctrl+u": m.viewport.ScrollUp(m.viewport.Height / 2) case "ctrl+d": m.viewport.ScrollDown(m.viewport.Height / 2) case "j": if m.inChoiceMode { return m, m.navigateList(tea.KeyDown) } case "k": if m.inChoiceMode { return m, m.navigateList(tea.KeyUp) } } // Default: send to text input var cmd tea.Cmd m.input, cmd = m.input.Update(msg) return m, cmd } return m, nil } func (m *model) handleEnter() (tea.Model, tea.Cmd) { // Handle choice selection if m.inChoiceMode { if _, ok := m.list.SelectedItem().(item); ok { if m.choiceType == "session" { idx := m.list.Index() if idx >= 0 && idx < len(m.sessionIDs) { selectedID := m.sessionIDs[idx] m.inChoiceMode = false m.choicePrompt = "" m.choiceOptionID = "" // Don't reset choiceType/sessionIDs yet or it might race, but actually we are done. m.dirty = true m.refresh() return m, func() tea.Msg { m.agent.Input <- &api.SessionPickerResponse{SessionID: selectedID} return nil } } } else { choice := m.list.Index() + 1 m.inChoiceMode = false m.choicePrompt = "" m.choiceOptionID = "" m.dirty = true m.refresh() return m, func() tea.Msg { m.agent.Input <- &api.UserChoiceResponse{Choice: choice} return nil } } } return m, nil } value := strings.TrimSpace(m.input.Value()) if value == "" { return m, nil } // Add user message m.messages = append(m.messages, &api.Message{ Source: api.MessageSourceUser, Type: api.MessageTypeText, Payload: value, Timestamp: time.Now(), }) m.input.Reset() m.dirty = true m.refresh() m.viewport.GotoBottom() // Intercept "sessions" command if value == "sessions" { return m, m.fetchSessions } m.thinkStart = time.Now() return m, func() tea.Msg { m.agent.Input <- &api.UserInputResponse{Query: value} return nil } } func (m *model) handleAgentMsg(msg *api.Message) (tea.Model, tea.Cmd) { session := m.agent.GetSession() m.messages = session.AllMessages() m.dirty = true // Check if we're entering choice mode - use the incoming message directly // to avoid race conditions where the message isn't yet in AllMessages() if msg.Type == api.MessageTypeUserChoiceRequest { if req, ok := msg.Payload.(*api.UserChoiceRequest); ok { items := make([]list.Item, len(req.Options)) for i, opt := range req.Options { items[i] = item(opt.Label) } m.list.SetItems(items) m.list.Select(0) m.inChoiceMode = true m.choicePrompt = req.Prompt m.choiceOptionID = msg.ID m.choiceType = "confirm" } } else if msg.Type == api.MessageTypeSessionPickerRequest { if req, ok := msg.Payload.(*api.SessionPickerRequest); ok { items := make([]list.Item, len(req.Sessions)) ids := make([]string, len(req.Sessions)) for i, s := range req.Sessions { label := fmt.Sprintf("%s (%s) • %d msgs", s.ID, s.ModelID, s.MessageCount) if s.Name != "" { label = fmt.Sprintf("%s (%s) • %s • %d msgs", s.Name, s.ModelID, s.ID, s.MessageCount) } items[i] = item(label) ids[i] = s.ID } m.list.SetItems(items) m.list.Select(0) m.inChoiceMode = true m.choicePrompt = "Select a session to resume" m.choiceOptionID = msg.ID m.choiceType = "session" m.sessionIDs = ids } } else if session.AgentState == api.AgentStateDone || session.AgentState == api.AgentStateExited { // Clear choice mode if we're done or exited m.inChoiceMode = false m.choicePrompt = "" m.choiceOptionID = "" } m.refresh() m.viewport.GotoBottom() if session.AgentState == api.AgentStateRunning || session.AgentState == api.AgentStateInitializing { return m, m.spinner.Tick } return m, nil } func (m *model) refresh() { if !m.dirty { return } m.viewport.SetContent(m.renderMessages()) m.dirty = false } func (m model) renderMessages() string { var sb strings.Builder if len(m.messages) == 0 { sb.WriteString(fmt.Sprintf("\n%s\n\n%s\n%s\n", primaryText.Render(logo), mutedStyle.PaddingLeft(1).Render("Your AI-powered Kubernetes assistant"), dimStyle.PaddingLeft(1).Render("Type a message to get started"))) } else { width := min(m.viewport.Width-6, 90) if width < 40 { width = 40 } renderer, err := m.cache.getRenderer(width) if err != nil { return "Error rendering messages" } for _, msg := range m.messages { if s := m.renderMessage(msg, renderer, width); s != "" { sb.WriteString(s) } } } // Render choice picker inline at the end of messages if m.inChoiceMode { sb.WriteString("\n") sb.WriteString(warnText.Render("? " + m.choicePrompt)) sb.WriteString("\n\n") sb.WriteString(m.list.View()) sb.WriteString("\n") } return sb.String() } func (m model) renderMessage(msg *api.Message, r *glamour.TermRenderer, w int) string { // Skip certain message types if msg.Type == api.MessageTypeUserInputRequest { if p, ok := msg.Payload.(string); ok && p == ">>>" { return "" } } if msg.Type == api.MessageTypeToolCallResponse { return "" } // Skip choice requests - they're rendered in the input area instead if msg.Type == api.MessageTypeUserChoiceRequest || msg.Type == api.MessageTypeSessionPickerRequest { return "" } // Check cache (except tool calls which show status) if msg.ID != "" && msg.Type != api.MessageTypeToolCallRequest { if cached, ok := m.cache.get(msg.ID); ok { return cached } } var result string switch msg.Type { case api.MessageTypeToolCallRequest: result = m.renderToolCall(msg, w) case api.MessageTypeError: result = m.renderError(msg, w) default: result = m.renderTextMsg(msg, r, w) } // Cache result if msg.ID != "" && result != "" && msg.Type != api.MessageTypeToolCallRequest { m.cache.set(msg.ID, result) } return result } func (m model) renderTextMsg(msg *api.Message, r *glamour.TermRenderer, w int) string { payload, ok := msg.Payload.(string) if !ok { return "" } ts := "" if !msg.Timestamp.IsZero() { ts = dimStyle.Italic(true).Render(" " + msg.Timestamp.Format("15:04")) } switch msg.Source { case api.MessageSourceUser: label := primaryText.Render("You") + ts content := textStyle.Width(w).Render(payload) return userMsg.Width(w+2).Render(label+"\n"+content) + "\n" case api.MessageSourceModel, api.MessageSourceAgent: label := successText.Render("kubectl-ai") + ts rendered, _ := r.Render(payload) return agentMsg.Width(w+2).Render(label+"\n"+strings.TrimSpace(rendered)) + "\n" } return "" } func (m model) renderToolCall(msg *api.Message, w int) string { payload, ok := msg.Payload.(string) if !ok { return "" } content := successText.Render("⚡ Running") + "\n" + codeStyle.Render(payload) return toolBox.Width(w).Render(content) + "\n" } func (m model) renderError(msg *api.Message, w int) string { payload, ok := msg.Payload.(string) if !ok { return "" } content := errorText.Render("✗ Error") + "\n" + errorText.Render(payload) return errorBox.Width(w).Render(content) + "\n" } func (m model) View() string { if m.quitting { return mutedStyle.Padding(1).Render("Goodbye!") } session := m.agent.GetSession() return lipgloss.JoinVertical(lipgloss.Left, m.viewStatus(session), m.viewDivider(), lipgloss.NewStyle().PaddingLeft(1).Render(m.viewport.View()), m.viewDivider(), m.viewInput(session.AgentState), m.viewHelp(session.AgentState), ) } func (m model) viewStatus(session *api.Session) string { sep := dimStyle.Render(" | ") name := session.Name if name == "" { name = session.ID } left := primaryText.Render("kubectl-ai") + sep + mutedStyle.Render(name) + sep + m.viewState(session.AgentState) model := session.ModelID if model == "" { model = "unknown" } right := lipgloss.NewStyle().Foreground(colorSecondary).Render(model) gap := m.width - lipgloss.Width(left) - lipgloss.Width(right) - 2 if gap < 0 { gap = 0 } return statusBar.Width(m.width).Render(" " + left + strings.Repeat(" ", gap) + right + " ") } func (m model) viewState(state api.AgentState) string { states := map[api.AgentState]struct { icon, text string style lipgloss.Style }{ api.AgentStateRunning: {"●", "Running", successText}, api.AgentStateInitializing: {"", "Initializing...", mutedStyle}, api.AgentStateWaitingForInput: {"●", "Ready", successText}, api.AgentStateIdle: {"○", "Idle", mutedStyle}, api.AgentStateDone: {"✓", "Done", successText}, api.AgentStateExited: {"○", "Exited", mutedStyle}, } if s, ok := states[state]; ok { txt := s.style.Render(s.icon + " " + s.text) if state == api.AgentStateRunning && !m.thinkStart.IsZero() { txt += mutedStyle.Render(" " + formatDuration(time.Since(m.thinkStart))) } return txt } return mutedStyle.Render(string(state)) } func (m model) viewDivider() string { return dimStyle.Render(strings.Repeat("─", m.width)) } func (m model) viewInput(state api.AgentState) string { // Show dimmed input hint when in choice mode (picker is inline above) if m.inChoiceMode { content := mutedStyle.Render("Use ↑/↓ to navigate, Enter to select") return lipgloss.NewStyle().Padding(0, 1).Render(inputBoxDim.Width(m.width - 4).Render(content)) } // Show spinner or input if state == api.AgentStateRunning || state == api.AgentStateInitializing { elapsed := "" if !m.thinkStart.IsZero() { elapsed = " " + formatDuration(time.Since(m.thinkStart)) } content := primaryText.Render(m.spinner.View()+" Thinking...") + mutedStyle.Render(elapsed) return lipgloss.NewStyle().Padding(0, 1).Render(inputBoxDim.Width(m.width - 4).Render(content)) } return lipgloss.NewStyle().Padding(0, 1).Render(inputBox.Width(m.width - 4).Render(m.input.View())) } func (m model) viewHelp(state api.AgentState) string { var hints []string if m.inChoiceMode { hints = []string{"↑/↓: navigate", "Enter: select", "Ctrl+C: quit"} } else if state == api.AgentStateRunning { hints = []string{"Ctrl+C: cancel"} } else { hints = []string{"Enter: send", "Esc: clear", "Ctrl+C: quit"} if m.viewport.TotalLineCount() > m.viewport.Height { hints = append(hints, "↑/↓: scroll") } } return dimStyle.Padding(0, 2, 1, 2).Render(strings.Join(hints, " • ")) } func formatDuration(d time.Duration) string { switch { case d < time.Minute: return fmt.Sprintf("%ds", int(d.Seconds())) case d < time.Hour: return fmt.Sprintf("%dm%ds", int(d.Minutes()), int(d.Seconds())%60) default: return fmt.Sprintf("%dh%dm", int(d.Hours()), int(d.Minutes())%60) } }