wehub/vxcontrol--pentagi
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:13 +08:00
commit 0878425be3
1160 changed files with 491311 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.dockerignore
+9
View File
@@ -0,0 +1,9 @@
frontend/coverage
frontend/dist
frontend/node_modules
frontend/ssl
**/*.log
**/*.env
**/.DS_Store
**/Thumbs.db
.env.example
+388
View File
@@ -0,0 +1,388 @@
# PentAGI Environment Variables
## For communication with PentAGI Cloud API
INSTALLATION_ID=
LICENSE_KEY=
## Allow to interact with user while executing tasks
ASK_USER=
## LLM Providers
OPEN_AI_KEY=
OPEN_AI_SERVER_URL=https://api.openai.com/v1
ANTHROPIC_API_KEY=
ANTHROPIC_SERVER_URL=https://api.anthropic.com/v1
## Google AI (Gemini) LLM provider
GEMINI_API_KEY=
GEMINI_SERVER_URL=https://generativelanguage.googleapis.com
## AWS Bedrock LLM provider
BEDROCK_REGION=us-east-1
BEDROCK_DEFAULT_AUTH=
BEDROCK_BEARER_TOKEN=
BEDROCK_ACCESS_KEY_ID=
BEDROCK_SECRET_ACCESS_KEY=
BEDROCK_SESSION_TOKEN=
BEDROCK_SERVER_URL=
## DeepSeek LLM provider
DEEPSEEK_API_KEY=
DEEPSEEK_SERVER_URL=https://api.deepseek.com
DEEPSEEK_PROVIDER=
## GLM (Zhipu AI) LLM provider
GLM_API_KEY=
GLM_SERVER_URL=https://api.z.ai/api/paas/v4
GLM_PROVIDER=
## Kimi (Moonshot) LLM provider
KIMI_API_KEY=
KIMI_SERVER_URL=https://api.moonshot.ai/v1
KIMI_PROVIDER=
## Qwen (Alibaba Cloud DashScope) LLM provider
QWEN_API_KEY=
QWEN_SERVER_URL=https://dashscope-us.aliyuncs.com/compatible-mode/v1
QWEN_PROVIDER=
## Custom LLM provider
LLM_SERVER_URL=
LLM_SERVER_KEY=
LLM_SERVER_MODEL=
LLM_SERVER_PROVIDER=
LLM_SERVER_CONFIG_PATH=
LLM_SERVER_LEGACY_REASONING=
LLM_SERVER_PRESERVE_REASONING=
## Ollama LLM provider (Local Server or Cloud)
# Local: http://ollama-server:11434, Cloud: https://ollama.com
OLLAMA_SERVER_URL=
# Required for Ollama Cloud (https://ollama.com/settings/keys), leave empty for local
OLLAMA_SERVER_API_KEY=
OLLAMA_SERVER_MODEL=
OLLAMA_SERVER_CONFIG_PATH=
OLLAMA_SERVER_PULL_MODELS_TIMEOUT=
OLLAMA_SERVER_PULL_MODELS_ENABLED=
OLLAMA_SERVER_LOAD_MODELS_ENABLED=
## Embedding
EMBEDDING_URL=
EMBEDDING_KEY=
EMBEDDING_MODEL=
EMBEDDING_PROVIDER=
EMBEDDING_BATCH_SIZE=
EMBEDDING_MAX_TEXT_BYTES=
EMBEDDING_STRIP_NEW_LINES=
## Summarizer
SUMMARIZER_PRESERVE_LAST=
SUMMARIZER_USE_QA=
SUMMARIZER_SUM_MSG_HUMAN_IN_QA=
SUMMARIZER_LAST_SEC_BYTES=
SUMMARIZER_MAX_BP_BYTES=
SUMMARIZER_MAX_QA_SECTIONS=
SUMMARIZER_MAX_QA_BYTES=
SUMMARIZER_KEEP_QA_SECTIONS=
## Assistant
ASSISTANT_USE_AGENTS=
ASSISTANT_SUMMARIZER_PRESERVE_LAST=
ASSISTANT_SUMMARIZER_LAST_SEC_BYTES=
ASSISTANT_SUMMARIZER_MAX_BP_BYTES=
ASSISTANT_SUMMARIZER_MAX_QA_SECTIONS=
ASSISTANT_SUMMARIZER_MAX_QA_BYTES=
ASSISTANT_SUMMARIZER_KEEP_QA_SECTIONS=
## Execution Monitor Detector
EXECUTION_MONITOR_ENABLED=
EXECUTION_MONITOR_SAME_TOOL_LIMIT=
EXECUTION_MONITOR_TOTAL_TOOL_LIMIT=
## Agent execution tool calls limit
MAX_GENERAL_AGENT_TOOL_CALLS=
MAX_LIMITED_AGENT_TOOL_CALLS=
## Agent planning step for pentester, coder, installer
AGENT_PLANNING_STEP_ENABLED=
## HTTP proxy to use it in isolation environment
PROXY_URL=
## SSL/TLS Certificate Configuration
EXTERNAL_SSL_CA_PATH=
EXTERNAL_SSL_INSECURE=
## HTTP client timeout in seconds for external API calls (LLM providers, search tools, etc.)
## Default: 600 (10 minutes). Set to 0 to use the default.
HTTP_CLIENT_TIMEOUT=
## Default terminal tool timeout in seconds applied when an agent requests timeout=0 or a negative value.
## Default: 1200 (20 minutes). Range: 1–10800 (up to 3 hours). Values <= 0 or above 10800 are clamped to 10800 (agents always get a finite timeout).
TERMINAL_TOOL_TIMEOUT=
## Scraper URLs and settings
## For Docker (default):
SCRAPER_PUBLIC_URL=
SCRAPER_PRIVATE_URL=https://someuser:somepass@scraper/
## For Podman rootless, use: SCRAPER_PRIVATE_URL=http://someuser:somepass@scraper:3000/
## See README.md "Running PentAGI with Podman" section for details
LOCAL_SCRAPER_USERNAME=someuser
LOCAL_SCRAPER_PASSWORD=somepass
LOCAL_SCRAPER_MAX_CONCURRENT_SESSIONS=10
## PentAGI server settings (docker-compose.yml)
PENTAGI_LISTEN_IP=
PENTAGI_LISTEN_PORT=
PENTAGI_DATA_DIR=
PENTAGI_SSL_DIR=
PENTAGI_OLLAMA_DIR=
PENTAGI_DOCKER_SOCKET=
PENTAGI_DOCKER_CERT_PATH=
PENTAGI_LLM_SERVER_CONFIG_PATH=
PENTAGI_OLLAMA_SERVER_CONFIG_PATH=
## PentAGI security settings
PUBLIC_URL=https://localhost:8443
CORS_ORIGINS=https://localhost:8443
COOKIE_SIGNING_SALT=salt # change this to improve security
## PentAGI internal server settings (inside the container)
STATIC_DIR=
STATIC_URL=
SERVER_PORT=8443
SERVER_HOST=0.0.0.0
SERVER_SSL_CRT=
SERVER_SSL_KEY=
SERVER_USE_SSL=true
## OAuth google
OAUTH_GOOGLE_CLIENT_ID=
OAUTH_GOOGLE_CLIENT_SECRET=
## OAuth github
OAUTH_GITHUB_CLIENT_ID=
OAUTH_GITHUB_CLIENT_SECRET=
## DuckDuckGo search engine
DUCKDUCKGO_ENABLED=
DUCKDUCKGO_REGION=
DUCKDUCKGO_SAFESEARCH=
DUCKDUCKGO_TIME_RANGE=
## Sploitus search engine API
SPLOITUS_ENABLED=
## Google search engine API
GOOGLE_API_KEY=
GOOGLE_CX_KEY=
GOOGLE_LR_KEY=
## Traversaal search engine API
TRAVERSAAL_API_KEY=
## Tavily search engine API
TAVILY_API_KEY=
## Perplexity search engine API
PERPLEXITY_API_KEY=
PERPLEXITY_MODEL=
PERPLEXITY_CONTEXT_SIZE=
## SEARXNG search engine API
SEARXNG_URL=
SEARXNG_CATEGORIES=general
SEARXNG_LANGUAGE=
SEARXNG_SAFESEARCH=0
SEARXNG_TIME_RANGE=
SEARXNG_TIMEOUT=
## Langfuse observability settings
LANGFUSE_BASE_URL=
LANGFUSE_PROJECT_ID=
LANGFUSE_PUBLIC_KEY=
LANGFUSE_SECRET_KEY=
## OpenTelemetry observability settings
OTEL_HOST=
## Docker client settings to run primary terminal container
DOCKER_HOST=
DOCKER_TLS_VERIFY=
DOCKER_CERT_PATH=
## Docker settings inside primary terminal container
DOCKER_INSIDE=true # enable to use docker socket
DOCKER_NET_ADMIN=true # enable to use net_admin capability
DOCKER_SOCKET=/var/run/docker.sock # path on host machine
DOCKER_NETWORK=
DOCKER_WORK_DIR=
DOCKER_PUBLIC_IP=0.0.0.0 # public ip of host machine
DOCKER_DEFAULT_IMAGE=
DOCKER_DEFAULT_IMAGE_FOR_PENTEST=
# Postgres (pgvector) settings
PENTAGI_POSTGRES_USER=postgres
PENTAGI_POSTGRES_PASSWORD=postgres # change this to improve security
PENTAGI_POSTGRES_DB=pentagidb
## Postgres (pgvector) connection pool settings
DATABASE_MAX_OPEN_CONNS=
DATABASE_MAX_IDLE_CONNS=
DATABASE_VECTOR_MAX_CONNS=
## Graphiti knowledge graph settings
## Set GRAPHITI_ENABLED=true and GRAPHITI_URL=http://graphiti:8000 to enable embedded Graphiti
GRAPHITI_ENABLED=false
GRAPHITI_TIMEOUT=30
GRAPHITI_URL=
GRAPHITI_MODEL_NAME=
# Neo4j settings (used by Graphiti stack)
NEO4J_USER=neo4j
NEO4J_DATABASE=neo4j
NEO4J_PASSWORD=devpassword # change this to improve security
NEO4J_URI=bolt://neo4j:7687
## PentAGI image settings
PENTAGI_IMAGE=
## Scraper network settings
## Default ports: SCRAPER_LISTEN_IP=127.0.0.1, SCRAPER_LISTEN_PORT=9443
## Note: These settings don't need to change for Podman rootless
SCRAPER_LISTEN_IP=
SCRAPER_LISTEN_PORT=
## Postgres network settings
PGVECTOR_LISTEN_IP=
PGVECTOR_LISTEN_PORT=
## Postgres Exporter network settings
POSTGRES_EXPORTER_LISTEN_IP=
POSTGRES_EXPORTER_LISTEN_PORT=
# Langfuse Environment Variables
## Langfuse server settings
LANGFUSE_LISTEN_IP=
LANGFUSE_LISTEN_PORT=
LANGFUSE_NEXTAUTH_URL=
## Langfuse Postgres
LANGFUSE_POSTGRES_USER=postgres
LANGFUSE_POSTGRES_PASSWORD=postgres # change this to improve security
LANGFUSE_POSTGRES_DB=langfuse
## Langfuse Clickhouse
LANGFUSE_CLICKHOUSE_USER=clickhouse
LANGFUSE_CLICKHOUSE_PASSWORD=clickhouse # change this to improve security
LANGFUSE_CLICKHOUSE_URL=http://langfuse-clickhouse:8123
LANGFUSE_CLICKHOUSE_MIGRATION_URL=clickhouse://langfuse-clickhouse:9000
LANGFUSE_CLICKHOUSE_CLUSTER_ENABLED=false
## Langfuse S3
LANGFUSE_S3_BUCKET=langfuse
LANGFUSE_S3_REGION=auto
LANGFUSE_S3_ACCESS_KEY_ID=accesskey # change this to improve security
LANGFUSE_S3_SECRET_ACCESS_KEY=secretkey # change this to improve security
LANGFUSE_S3_ENDPOINT=http://langfuse-minio:9000
LANGFUSE_S3_FORCE_PATH_STYLE=true
LANGFUSE_S3_EVENT_UPLOAD_PREFIX=events/
LANGFUSE_S3_MEDIA_UPLOAD_PREFIX=media/
LANGFUSE_S3_BATCH_EXPORT_ENABLED=true
## Langfuse Redis
LANGFUSE_REDIS_HOST=langfuse-redis
LANGFUSE_REDIS_PORT=6379
LANGFUSE_REDIS_AUTH=redispassword # change this to improve security
LANGFUSE_REDIS_TLS_ENABLED=false
LANGFUSE_REDIS_TLS_CA=
LANGFUSE_REDIS_TLS_CERT=
LANGFUSE_REDIS_TLS_KEY=
## Langfuse web app security settings
LANGFUSE_SALT=salt # change this to improve security
LANGFUSE_ENCRYPTION_KEY=0000000000000000000000000000000000000000000000000000000000000000 # change this to improve security
## Langfuse web app nextauth settings
LANGFUSE_NEXTAUTH_URL=http://localhost:4000
LANGFUSE_NEXTAUTH_SECRET=secret # change this to improve security
## Langfuse extra settings
LANGFUSE_ENABLE_EXPERIMENTAL_FEATURES=true
LANGFUSE_TELEMETRY_ENABLED=false
LANGFUSE_LOG_LEVEL=info
## Langfuse init settings
LANGFUSE_INIT_ORG_ID=ocm47619l0000872mcd2dlbqwb
LANGFUSE_INIT_ORG_NAME=PentAGI Org
LANGFUSE_INIT_PROJECT_ID=cm47619l0000872mcd2dlbqwb
LANGFUSE_INIT_PROJECT_NAME=PentAGI
LANGFUSE_INIT_PROJECT_PUBLIC_KEY=pk-lf-00000000-0000-0000-0000-000000000000 # change this to improve security
LANGFUSE_INIT_PROJECT_SECRET_KEY=sk-lf-00000000-0000-0000-0000-000000000000 # change this to improve security
LANGFUSE_INIT_USER_EMAIL=admin@pentagi.com
LANGFUSE_INIT_USER_NAME=admin
LANGFUSE_INIT_USER_PASSWORD=password # change this to improve security
## Langfuse SDK sync settings
LANGFUSE_SDK_CI_SYNC_PROCESSING_ENABLED=false
LANGFUSE_READ_FROM_POSTGRES_ONLY=false
LANGFUSE_READ_FROM_CLICKHOUSE_ONLY=true
LANGFUSE_RETURN_FROM_CLICKHOUSE=true
## Langfuse ingestion tuning
LANGFUSE_INGESTION_QUEUE_DELAY_MS=
LANGFUSE_INGESTION_CLICKHOUSE_WRITE_INTERVAL_MS=
LANGFUSE_INGESTION_CLICKHOUSE_WRITE_BATCH_SIZE=
LANGFUSE_INGESTION_CLICKHOUSE_MAX_ATTEMPTS=
## Langfuse email
LANGFUSE_EMAIL_FROM_ADDRESS=
LANGFUSE_SMTP_CONNECTION_URL=
## Langfuse optional Azure blob
LANGFUSE_USE_AZURE_BLOB=false
## Langfuse license settings
LANGFUSE_EE_LICENSE_KEY=
## Langfuse OpenTelemetry settings
LANGFUSE_OTEL_EXPORTER_OTLP_ENDPOINT=
LANGFUSE_OTEL_SERVICE_NAME=
## Langfuse custom oauth2 settings
LANGFUSE_AUTH_CUSTOM_CLIENT_ID=
LANGFUSE_AUTH_CUSTOM_CLIENT_SECRET=
LANGFUSE_AUTH_CUSTOM_ISSUER=
LANGFUSE_AUTH_CUSTOM_NAME=PentAGI
LANGFUSE_AUTH_CUSTOM_SCOPE=openid email profile
LANGFUSE_AUTH_CUSTOM_CLIENT_AUTH_METHOD=client_secret_post
LANGFUSE_AUTH_CUSTOM_ALLOW_ACCOUNT_LINKING=true
## Langfuse auth settings
LANGFUSE_AUTH_DISABLE_SIGNUP=false # disable signup if PentAGI OAuth2 is used
LANGFUSE_AUTH_SESSION_MAX_AGE=240
## Langfuse allowed organization creators
LANGFUSE_ALLOWED_ORGANIZATION_CREATORS=admin@pentagi.com
## Langfuse default settings for new users
LANGFUSE_DEFAULT_ORG_ID=ocm47619l0000872mcd2dlbqwb
LANGFUSE_DEFAULT_PROJECT_ID=cm47619l0000872mcd2dlbqwb
LANGFUSE_DEFAULT_ORG_ROLE=VIEWER
LANGFUSE_DEFAULT_PROJECT_ROLE=VIEWER
# Observability Environment Variables
## Observability server settings
GRAFANA_LISTEN_IP=
GRAFANA_LISTEN_PORT=
## OpenTelemetry server settings
OTEL_GRPC_LISTEN_IP=
OTEL_GRPC_LISTEN_PORT=
OTEL_HTTP_LISTEN_IP=
OTEL_HTTP_LISTEN_PORT=
.github/ISSUE_TEMPLATE/1-bug-report.yml
+126
View File
@@ -0,0 +1,126 @@
name: "\U0001F41B Bug report"
description: "Report a bug in PentAGI"
title: "[Bug]: "
labels: ["bug"]
assignees:
- asdek
body:
- type: markdown
attributes:
value: |
Thanks for taking the time to fill out this bug report! Please provide as much information as possible to help us diagnose and fix the issue.
- type: dropdown
id: component
attributes:
label: Affected Component
description: Which component of PentAGI is affected by this bug?
multiple: true
options:
- Core Services (Frontend UI/Backend API)
- AI Agents (Researcher/Developer/...)
- Security Tools Integration
- Memory System (Vector Store/Knowledge Base)
- Monitoring Stack Integration (Grafana/OpenTelemetry)
- Analytics Platform Integration (Langfuse)
- External Integrations (LLM/Search APIs)
- Documentation and User Experience
- Other (please specify in the description)
validations:
required: true
- type: textarea
attributes:
label: Describe the bug
description: Please provide a clear and concise description of the bug, including expected and actual behavior.
placeholder: |
What happened:
- Actual behavior: When executing a penetration test against [target], the AI agent [behavior]
What should happen:
- Expected behavior: The system should [expected outcome]
Additional context:
- Task/Flow ID (if applicable): [ID from UI]
- Error messages: [any error messages from logs/UI]
validations:
required: true
- type: textarea
attributes:
label: Steps to Reproduce
description: Please provide detailed steps to reproduce the bug.
placeholder: |
1. Access PentAGI Web UI at [relative URL]
2. Start a new flow with parameters [...] or prompt [...]
3. Configure target system as [...]
4. Observe AI agent behavior in [...] or log from Langfuse
5. Error occurs when [...] or screenshot/export logs from Grafana
validations:
required: true
- type: textarea
attributes:
label: System Configuration
description: Please provide details about your setup
placeholder: |
PentAGI Version: [e.g., latest from Docker Hub]
Deployment Type:
- [ ] Docker Compose
- [ ] Custom Deployment
Environment:
- Docker Version: [output of `docker --version`]
- Docker Compose Version: [output of `docker compose version`]
- Host OS: [e.g., Ubuntu 22.04, macOS 14.0]
- Available Resources:
- RAM: [e.g., 8GB]
- CPU: [e.g., 4 cores]
- Disk Space: [e.g., 50GB free]
Enabled Features:
- [ ] Langfuse Analytics
- [ ] Grafana Monitoring
- [ ] Custom LLM Server
Active Integrations:
- LLM Provider: [OpenAI/Anthropic/Custom]
- Search Systems: [Google/DuckDuckGo/Tavily/Traversaal/Perplexity]
validations:
required: true
- type: textarea
attributes:
label: Logs and Artifacts
description: |
Please provide relevant logs and artifacts. You can find logs using:
- Docker logs: `docker logs pentagi`
- Grafana dashboards (if enabled)
- Langfuse traces (if enabled)
- Browser console logs (for UI issues)
placeholder: |
```
Paste logs here
```
For large logs, please use GitHub Gist and provide the link.
validations:
required: false
- type: textarea
attributes:
label: Screenshots or Recordings
description: |
If applicable, add screenshots or recordings to help explain your problem.
- For UI issues: Browser screenshots/recordings
- For agent behavior: Langfuse trace screenshots
- For monitoring: Grafana dashboard screenshots
placeholder: Drag and drop images/videos here, or paste links to external storage.
validations:
required: false
- type: checkboxes
id: verification
attributes:
label: Verification
description: Please verify the following before submitting
options:
- label: I have checked that this issue hasn't been already reported
- label: I have provided all relevant configuration files (with sensitive data removed)
- label: I have included relevant logs and error messages
- label: I am running the latest version of PentAGI
validations:
required: true
.github/ISSUE_TEMPLATE/2-enhancement.yml
+115
View File
@@ -0,0 +1,115 @@
name: "\U0001F680 Enhancement"
description: "Suggest an enhancement for PentAGI"
title: "[Enhancement]: "
labels: ["enhancement"]
assignees:
- asdek
body:
- type: markdown
attributes:
value: |
Thank you for suggesting an enhancement to make PentAGI better! Please provide as much detail as possible to help us understand your suggestion.
- type: dropdown
id: component
attributes:
label: Target Component
description: Which component of PentAGI would this enhancement affect?
multiple: true
options:
- Core Services (Frontend UI/Backend API)
- AI Agents (Researcher/Developer/Executor)
- Security Tools Integration
- Memory System (Vector Store/Knowledge Base)
- Monitoring Stack (Grafana/OpenTelemetry)
- Analytics Platform (Langfuse)
- External Integrations (LLM/Search APIs)
- Documentation and User Experience
validations:
required: true
- type: textarea
attributes:
label: Enhancement Description
description: Please describe the enhancement you would like to see.
placeholder: |
Problem Statement:
- Current Limitation: [describe what's currently missing or could be improved]
- Use Case: [describe how you use PentAGI and why this enhancement would help]
Proposed Solution:
- Feature Description: [detailed description of the enhancement]
- Expected Benefits: [how this would improve PentAGI]
Example Scenario:
[Provide a concrete example of how this enhancement would be used]
validations:
required: true
- type: textarea
attributes:
label: Technical Details
description: If you have technical suggestions for implementation, please share them.
placeholder: |
Implementation Approach:
- Architecture Changes: [any changes needed to current architecture]
- New Components: [any new services or integrations needed]
- Dependencies: [new tools or libraries required]
Integration Points:
- AI Agents: [how it affects agent behavior]
- Memory System: [data storage requirements]
- Monitoring: [new metrics or traces needed]
Security Considerations:
- [Any security implications to consider]
validations:
required: false
- type: textarea
attributes:
label: Designs and Mockups
description: |
If applicable, provide mockups, diagrams, or examples to illustrate your enhancement.
- For UI changes: wireframes or mockups
- For architecture changes: system diagrams
- For agent behavior: sequence diagrams
placeholder: |
Drag and drop images here, or provide links to external design tools.
For complex diagrams, you can use Mermaid syntax:
```mermaid
sequenceDiagram
User->>PentAGI: Request
PentAGI->>NewComponent: Process
NewComponent->>User: Enhanced Response
```
validations:
required: false
- type: textarea
attributes:
label: Alternative Solutions
description: |
Please describe any alternative solutions or features you've considered.
placeholder: |
Alternative Approaches:
1. [First alternative approach]
- Pros: [benefits]
- Cons: [drawbacks]
2. [Second alternative approach]
- Pros: [benefits]
- Cons: [drawbacks]
Reason for Preferred Solution:
[Explain why your main proposal is better than these alternatives]
validations:
required: false
- type: checkboxes
id: verification
attributes:
label: Verification
description: Please verify the following before submitting
options:
- label: I have checked that this enhancement hasn't been already proposed
- label: This enhancement aligns with PentAGI's goal of autonomous penetration testing
- label: I have considered the security implications of this enhancement
- label: I have provided clear use cases and benefits
validations:
required: true
.github/PULL_REQUEST_TEMPLATE.md
+120
View File
@@ -0,0 +1,120 @@
<!--
Thank you for your contribution to PentAGI! Please fill out this template completely to help us review your changes effectively.
Any PR that does not include enough information may be closed at maintainers' discretion.
-->
### Description of the Change
<!--
We must be able to understand the design of your change from this description. Please provide as much detail as possible.
-->
#### Problem
<!-- Describe the problem this PR addresses -->
#### Solution
<!-- Describe your solution and its key aspects -->
<!-- Enter any applicable Issue number(s) here that will be closed/resolved by this PR. -->
Closes #
### Type of Change
<!-- Mark with an `x` all options that apply -->
- [ ] 🐛 Bug fix (non-breaking change which fixes an issue)
- [ ] 🚀 New feature (non-breaking change which adds functionality)
- [ ] 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] 📚 Documentation update
- [ ] 🔧 Configuration change
- [ ] 🧪 Test update
- [ ] 🛡️ Security update
### Areas Affected
<!-- Mark with an `x` all components that are affected -->
- [ ] Core Services (Frontend UI/Backend API)
- [ ] AI Agents (Researcher/Developer/Executor)
- [ ] Security Tools Integration
- [ ] Memory System (Vector Store/Knowledge Base)
- [ ] Monitoring Stack (Grafana/OpenTelemetry)
- [ ] Analytics Platform (Langfuse)
- [ ] External Integrations (LLM/Search APIs)
- [ ] Documentation
- [ ] Infrastructure/DevOps
### Testing and Verification
<!--
Please describe the tests that you ran to verify your changes and provide instructions so we can reproduce.
-->
#### Test Configuration
```yaml
PentAGI Version:
Docker Version:
Host OS:
LLM Provider:
Enabled Features: [Langfuse/Grafana/etc]
```
#### Test Steps
1.
2.
3.
#### Test Results
<!-- Include relevant screenshots, logs, or test outputs -->
### Security Considerations
<!--
Describe any security implications of your changes.
For security-related changes, please note any new dependencies, changed permissions, etc.
-->
### Performance Impact
<!--
Describe any performance implications and testing done to verify acceptable performance.
Especially important for changes affecting AI agents, memory systems, or data processing.
-->
### Documentation Updates
<!-- Note any documentation changes required by this PR -->
- [ ] README.md updates
- [ ] API documentation updates
- [ ] Configuration documentation updates
- [ ] GraphQL schema updates
- [ ] Other: <!-- specify -->
### Deployment Notes
<!--
Describe any special considerations for deploying this change.
Include any new environment variables, configuration changes, or migration steps.
-->
### Checklist
<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
#### Code Quality
- [ ] My code follows the project's coding standards
- [ ] I have added/updated necessary documentation
- [ ] I have added tests to cover my changes
- [ ] All new and existing tests pass
- [ ] I have run `go fmt` and `go vet` (for Go code)
- [ ] I have run `pnpm run lint` (for TypeScript/JavaScript code)
#### Security
- [ ] I have considered security implications
- [ ] Changes maintain or improve the security model
- [ ] Sensitive information has been properly handled
#### Compatibility
- [ ] Changes are backward compatible
- [ ] Breaking changes are clearly marked and documented
- [ ] Dependencies are properly updated
#### Documentation
- [ ] Documentation is clear and complete
- [ ] Comments are added for non-obvious code
- [ ] API changes are documented
### Additional Notes
<!-- Any additional information that would be helpful for reviewers -->
.github/SAVED_REPLIES.md
+176
View File
@@ -0,0 +1,176 @@
# Saved Replies
These are standardized responses for the PentAGI Development Team to use when responding to Issues and Pull Requests. Using these templates helps maintain consistency in our communications and saves time.
Since GitHub currently does not support repository-wide saved replies, team members should maintain these individually. All responses are versioned for easier updates.
While these are templates, please customize them to fit the specific context and:
- Welcome new contributors
- Thank them for their contribution
- Provide context for your response
- Outline next steps
You can add these saved replies to [your personal GitHub account here](https://github.com/settings/replies).
## Issue Responses
### Issue: Already Fixed (v1)
```
Thank you for reporting this issue! This has been resolved in a recent release. Please update to the latest version (see our [releases page](https://github.com/vxcontrol/pentagi/releases)) and verify if the issue persists.
If you continue experiencing problems after updating, please:
1. Check your configuration against our documentation
2. Provide logs from both PentAGI and monitoring systems (Grafana/Langfuse)
3. Include details about your environment and enabled features
```
### Issue: Need More Information (v1)
```
Thank you for your report! To help us better understand and address your issue, please provide additional information:
1. PentAGI version and deployment method (Docker Compose/Custom)
2. Relevant logs from:
- Docker containers
- Grafana dashboards (if enabled)
- Langfuse traces (if enabled)
3. Steps to reproduce the issue
4. Expected vs actual behavior
Please update your issue using our bug report template for consistency.
```
### Issue: Cannot Reproduce (v1)
```
Thank you for reporting this issue! Unfortunately, I cannot reproduce the problem with the provided information. To help us investigate:
1. Verify you're using the latest version
2. Provide your complete environment configuration
3. Share relevant logs and monitoring data
4. Include step-by-step reproduction instructions
5. Specify which AI agents were involved (Researcher/Developer/Executor)
Please update your issue with these details so we can better assist you.
```
### Issue: Expected Behavior (v1)
```
Thank you for your report! This appears to be the expected behavior because:
[Explanation of why this is working as designed]
If you believe this behavior should be different, please:
1. Describe your use case in detail
2. Explain why the current behavior doesn't meet your needs
3. Suggest alternative behavior that would work better
We're always open to improving PentAGI's functionality.
```
### Issue: Missing Template (v1)
```
Thank you for reporting this! To help us process your issue efficiently, please use our issue templates:
- [Bug Report Template](https://github.com/vxcontrol/pentagi/blob/master/.github/ISSUE_TEMPLATE/1-bug-report.md) for problems
- [Enhancement Template](https://github.com/vxcontrol/pentagi/blob/master/.github/ISSUE_TEMPLATE/2-enhancement.md) for suggestions
Please edit your issue to include the template information. This helps ensure we have all necessary details to assist you.
```
### Issue: PR Welcome (v1)
```
Thank you for raising this issue! We welcome contributions from the community.
If you'd like to implement this yourself:
1. Check our [contribution guidelines](CONTRIBUTING.md)
2. Review the architecture documentation
3. Consider security implications (especially for AI agent modifications)
4. Include tests and documentation
5. Update monitoring/analytics as needed
Feel free to ask questions if you need guidance. We're here to help!
```
## PR Responses
### PR: Ready to Merge (v1)
```
Excellent work! This PR meets our quality standards and I'll proceed with merging it.
If you're interested in further contributions, check our:
- [Help Wanted Issues](https://github.com/vxcontrol/pentagi/labels/help-wanted)
- [Good First Issues](https://github.com/vxcontrol/pentagi/labels/good-first-issue)
Thank you for improving PentAGI!
```
### PR: Needs Work (v1)
```
Thank you for your contribution! A few items need attention before we can merge:
[List specific items that need addressing]
Common requirements:
- Tests for new functionality
- Documentation updates
- Security considerations
- Performance impact assessment
- Monitoring/analytics integration
Please update your PR addressing these points. Let us know if you need any clarification.
```
### PR: Missing Template (v1)
```
Thank you for your contribution! Please update your PR to use our [PR template](https://github.com/vxcontrol/pentagi/blob/master/.github/PULL_REQUEST_TEMPLATE.md).
The template helps ensure we have:
- Clear description of changes
- Testing information
- Security considerations
- Documentation updates
- Deployment notes
This helps us review your changes effectively.
```
### PR: Missing Issue (v1)
```
Thank you for your contribution! We require an associated issue for each PR to:
- Discuss approach before implementation
- Track related changes
- Maintain clear project history
Please:
1. [Create an issue](https://github.com/vxcontrol/pentagi/issues/new/choose)
2. Link it to this PR
3. Update the PR description with the issue reference
This helps us maintain good project organization.
```
### PR: Inactive (v1)
```
This PR has been inactive for a while. To keep our review process efficient:
1. If you're still working on this:
- Let us know your timeline
- Update with latest main branch
- Address any existing feedback
2. If you're no longer working on this:
- We can close it
- Someone else can pick it up
Please let us know your preference within the next week.
```
### General: Need Help (v1)
```
I need additional expertise on this. Pinging:
- @asdek for technical review
- @security-team for security implications
- @ai-team for AI agent behavior
- @infra-team for infrastructure changes
[Specific questions or concerns that need addressing]
```
.github/workflows/ci.yml
+206
View File
@@ -0,0 +1,206 @@
name: Docker build and push
on:
push:
branches:
- "**"
tags:
- "v[0-9]+.[0-9]+.[0-9]+"
workflow_dispatch:
jobs:
lint-and-test:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v6
with:
fetch-depth: 0
# Go setup and cache
- name: Set up Go
uses: actions/setup-go@v6
with:
go-version: '1.24'
cache: true
cache-dependency-path: backend/go.sum
# Cache Go dependencies
- name: Go Mod Cache
uses: actions/cache@v5
with:
path: |
~/.cache/go-build
~/go/pkg/mod
key: ${{ runner.os }}-go-${{ hashFiles('backend/go.sum') }}
restore-keys: |
${{ runner.os }}-go-
# pnpm setup
- name: Install pnpm
uses: pnpm/action-setup@v4
with:
package_json_file: frontend/package.json
# Node.js setup and cache
- name: Set up Node.js
uses: actions/setup-node@v6
with:
node-version: '23'
cache: 'pnpm'
cache-dependency-path: 'frontend/pnpm-lock.yaml'
# Frontend lint and test
- name: Frontend - Install dependencies
working-directory: frontend
run: pnpm install --frozen-lockfile
continue-on-error: true
- name: Frontend - Prettier
working-directory: frontend
run: pnpm run prettier
continue-on-error: true
- name: Frontend - Lint
working-directory: frontend
run: pnpm run lint
continue-on-error: true
- name: Frontend - Test
working-directory: frontend
run: pnpm run test
continue-on-error: true
# Backend lint and test
- name: Backend - Download dependencies
working-directory: backend
run: go mod download
continue-on-error: true
- name: Backend - Lint
uses: golangci/golangci-lint-action@v9
with:
version: latest
working-directory: backend
args: --timeout=5m --issues-exit-code=0
continue-on-error: true
- name: Backend - Test
working-directory: backend
run: go test ./... -v
continue-on-error: true
- name: Backend - Test Build
working-directory: backend
env:
CGO_ENABLED: 0
GO111MODULE: on
run: |
# Get version information
LATEST_TAG=$(git describe --tags --abbrev=0 2>/dev/null || echo "v0.0.0")
PACKAGE_VER=${LATEST_TAG#v}
CURRENT_COMMIT=$(git rev-parse HEAD)
TAG_COMMIT=$(git rev-list -n 1 "$LATEST_TAG" 2>/dev/null || echo "")
if [ "$CURRENT_COMMIT" != "$TAG_COMMIT" ]; then
PACKAGE_REV=$(git rev-parse --short HEAD)
else
PACKAGE_REV=""
fi
LDFLAGS="-X pentagi/pkg/version.PackageName=pentagi -X pentagi/pkg/version.PackageVer=${PACKAGE_VER} -X pentagi/pkg/version.PackageRev=${PACKAGE_REV}"
echo "Building with version: ${PACKAGE_VER}${PACKAGE_REV:+-$PACKAGE_REV}"
# Build for AMD64
GOOS=linux GOARCH=amd64 go build -trimpath -ldflags "$LDFLAGS" -o /tmp/pentagi-amd64 ./cmd/pentagi
echo "✓ Successfully built for linux/amd64"
# Build for ARM64
GOOS=linux GOARCH=arm64 go build -trimpath -ldflags "$LDFLAGS" -o /tmp/pentagi-arm64 ./cmd/pentagi
echo "✓ Successfully built for linux/arm64"
continue-on-error: true
docker-build:
needs: lint-and-test
if: github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/v')
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v6
with:
fetch-depth: 0
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Login to Docker Hub
uses: docker/login-action@v3
with:
username: ${{ secrets.DOCKERHUB_USERNAME }}
password: ${{ secrets.DOCKERHUB_TOKEN }}
# Extract version from tag (without 'v' prefix) and split into parts
- name: Extract version and revision
id: version
run: |
# Get latest tag version (without 'v' prefix)
LATEST_TAG=$(git describe --tags --abbrev=0 2>/dev/null || echo "v0.0.0")
VERSION=${LATEST_TAG#v}
# Get current commit hash
CURRENT_COMMIT=$(git rev-parse HEAD)
# Get commit hash of the latest tag
TAG_COMMIT=$(git rev-list -n 1 "$LATEST_TAG" 2>/dev/null || echo "")
# Set revision only if current commit differs from tag commit
if [ "$CURRENT_COMMIT" != "$TAG_COMMIT" ]; then
PACKAGE_REV=$(git rev-parse --short HEAD)
echo "revision=${PACKAGE_REV}" >> $GITHUB_OUTPUT
echo "is_release=false" >> $GITHUB_OUTPUT
echo "Building development version: ${VERSION}-${PACKAGE_REV}"
echo " Docker tags: latest only"
else
echo "revision=" >> $GITHUB_OUTPUT
echo "is_release=true" >> $GITHUB_OUTPUT
echo "Building release version: ${VERSION}"
# Split version into major.minor.patch for Docker tags (only for releases)
IFS='.' read -r major minor patch <<< "$VERSION"
echo "major=${major}" >> $GITHUB_OUTPUT
echo "minor=${major}.${minor}" >> $GITHUB_OUTPUT
echo "patch=${VERSION}" >> $GITHUB_OUTPUT
echo " Docker tags: latest, ${major}, ${major}.${minor}, ${VERSION}"
fi
echo "version=${VERSION}" >> $GITHUB_OUTPUT
- name: Generate Docker metadata
id: meta
uses: docker/metadata-action@v5
with:
images: vxcontrol/pentagi
tags: |
# For main branch - latest tag
type=raw,value=latest,enable=${{ github.ref == 'refs/heads/main' }}
# For release builds only - tag with major, minor and patch versions
type=raw,value=${{ steps.version.outputs.major }},enable=${{ steps.version.outputs.is_release == 'true' }}
type=raw,value=${{ steps.version.outputs.minor }},enable=${{ steps.version.outputs.is_release == 'true' }}
type=raw,value=${{ steps.version.outputs.patch }},enable=${{ steps.version.outputs.is_release == 'true' }}
- name: Build and push
uses: docker/build-push-action@v6
with:
context: .
platforms: linux/amd64,linux/arm64/v8
push: true
provenance: true
sbom: true
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha
cache-to: type=gha,mode=max
build-args: |
PACKAGE_VER=${{ steps.version.outputs.version }}
PACKAGE_REV=${{ steps.version.outputs.revision }}
.gitignore
+37
View File
@@ -0,0 +1,37 @@
.DS_Store
.env
.env.*
.state
!.env.example
!backend/cmd/installer/files/links/.env
backend/tmp
backend/build
frontend/coverage
frontend/dist
frontend/node_modules
frontend/ssl
node_modules
.cursorrules
.cursorignore
.cursor/
.agents/
skills-lock.json
build/*
data/*
.bak/*
!.gitkeep
.claude/
# IDE
.idea/
*.swp
*.swo
*~
# OS
Thumbs.db
.vscode/launch.json
Vendored
+164
View File
@@ -0,0 +1,164 @@
{
"version": "0.2.0",
"configurations": [
{
"type": "go",
"request": "launch",
"name": "Launch Go Backend",
"program": "${workspaceFolder}/backend/cmd/pentagi/main.go",
"envFile": "${workspaceFolder}/.env",
"env": {
"CORS_ORIGINS": "http://localhost:*,https://localhost:*",
"SERVER_PORT": "8080",
"SERVER_USE_SSL": "false",
"DATABASE_URL": "postgres://postgres:postgres@localhost:5432/pentagidb?sslmode=disable",
"PUBLIC_URL": "http://localhost:8080",
"STATIC_URL": "http://localhost:8000",
// Choose it instead of STATIC_URL to serve static files from a directory:
// "STATIC_DIR": "${workspaceFolder}/frontend/dist",
// Langfuse (optional) uncomment to enable
// "LANGFUSE_BASE_URL": "http://localhost:4000",
// Observability (optional) uncomment to enable
// "OTEL_HOST": "localhost:8148",
// Scraper (optional) uncomment to enable
// "SCRAPER_PRIVATE_URL": "https://someuser:somepass@localhost:9443/",
},
"cwd": "${workspaceFolder}",
"output": "${workspaceFolder}/build/__debug_bin_pentagi",
},
{
"type": "node",
"request": "launch",
"name": "Launch Frontend",
"runtimeExecutable": "pnpm",
"runtimeArgs": ["run", "dev"],
"env": {
"VITE_APP_LOG_LEVEL": "DEBUG",
"VITE_API_URL": "localhost:8080",
"VITE_USE_HTTPS": "false",
"VITE_PORT": "8000",
"VITE_HOST": "0.0.0.0",
},
"cwd": "${workspaceFolder}/frontend",
},
{
"type": "chrome",
"request": "launch",
"name": "Launch Browser",
"url": "http://localhost:8000",
"webRoot": "${workspaceFolder}/frontend/src",
},
{
"type": "go",
"request": "launch",
"name": "Launch Agents Tests",
"program": "${workspaceFolder}/backend/cmd/ctester/",
"envFile": "${workspaceFolder}/.env",
"env": {},
"args": [
// "-type", "openai",
// "-type", "anthropic",
// "-type", "gemini",
// "-type", "bedrock",
// "-type", "ollama",
// "-type", "deepseek",
// "-type", "glm",
// "-type", "kimi",
// "-type", "qwen",
"-config", "${workspaceFolder}/examples/configs/moonshot.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/deepseek.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/ollama-cloud.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/ollama-llama318b.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/ollama-llama318b-instruct.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/ollama-qwq32b-fp16-tc.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/ollama-qwen332b-fp16-tc.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/vllm-qwen3.5-27b-fp8.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/vllm-qwen332b-fp16.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/custom-openai.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/openrouter.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/deepinfra.provider.yml",
// "-config", "${workspaceFolder}/examples/configs/novita.provider.yml",
// "-report", "${workspaceFolder}/examples/tests/openai-report.md",
// "-report", "${workspaceFolder}/examples/tests/anthropic-report.md",
// "-report", "${workspaceFolder}/examples/tests/gemini-report.md",
// "-report", "${workspaceFolder}/examples/tests/bedrock-report.md",
// "-report", "${workspaceFolder}/examples/tests/ollama-cloud-report.md",
// "-report", "${workspaceFolder}/examples/tests/ollama-llama318b-report.md",
// "-report", "${workspaceFolder}/examples/tests/ollama-llama318b-instruct-report.md",
// "-report", "${workspaceFolder}/examples/tests/ollama-qwq-32b-fp16-tc-report.md",
// "-report", "${workspaceFolder}/examples/tests/ollama-qwen332b-fp16-tc-report.md",
// "-report", "${workspaceFolder}/examples/tests/vllm-qwen3.5-27b-fp8.report.md",
// "-report", "${workspaceFolder}/examples/tests/vllm-qwen332b-fp16-report.md",
"-report", "${workspaceFolder}/examples/tests/moonshot-report.md",
// "-report", "${workspaceFolder}/examples/tests/deepseek-report.md",
// "-report", "${workspaceFolder}/examples/tests/glm-report.md",
// "-report", "${workspaceFolder}/examples/tests/kimi-report.md",
// "-report", "${workspaceFolder}/examples/tests/qwen-report.md",
// "-report", "${workspaceFolder}/examples/tests/custom-openai-report.md",
// "-report", "${workspaceFolder}/examples/tests/openrouter-report.md",
// "-report", "${workspaceFolder}/examples/tests/deepinfra-report.md",
// "-report", "${workspaceFolder}/examples/tests/novita-report.md",
"-agents", "all",
"-groups", "all",
"-workers", "8",
"-verbose",
],
"cwd": "${workspaceFolder}",
"output": "${workspaceFolder}/build/__debug_bin_ctester",
},
{
"type": "go",
"request": "launch",
"name": "Launch Installer",
"program": "${workspaceFolder}/backend/cmd/installer/",
"args": ["-e", "${workspaceFolder}/.env_test.example"],
"cwd": "${workspaceFolder}",
"output": "${workspaceFolder}/build/__debug_bin_installer",
},
{
"type": "go",
"request": "attach",
"name": "Attach Installer",
"mode": "local",
"processId": "${command:PickProcess}",
"cwd": "${workspaceFolder}",
},
{
"type": "go",
"request": "launch",
"name": "Launch Tools Tests",
"program": "${workspaceFolder}/backend/cmd/ftester/",
"envFile": "${workspaceFolder}/.env",
"env": {
"LLM_SERVER_CONFIG_PATH": "${workspaceFolder}/examples/configs/openrouter.provider.yml",
"DATABASE_URL": "postgres://postgres:postgres@localhost:5432/pentagidb?sslmode=disable",
// Langfuse (optional) uncomment to enable
"LANGFUSE_BASE_URL": "http://localhost:4000",
// Observability (optional) uncomment to enable
"OTEL_HOST": "localhost:8148",
},
"args": [
"-flow", "0",
// "describe", "-verbose",
"perplexity", "-query", "how I can install burp suite on Kali Linux?", "-max_results", "10",
],
"cwd": "${workspaceFolder}",
"output": "${workspaceFolder}/build/__debug_bin_ftester",
},
{
"type": "go",
"request": "launch",
"name": "Launch Embedding Tests",
"program": "${workspaceFolder}/backend/cmd/etester/",
"envFile": "${workspaceFolder}/.env",
"env": {
"DATABASE_URL": "postgres://postgres:postgres@localhost:5432/pentagidb?sslmode=disable",
},
"args": [
"info", "-verbose",
],
"cwd": "${workspaceFolder}",
"output": "${workspaceFolder}/build/__debug_bin_etester",
},
]
}
.vscode/settings.json
Vendored
+34
View File
@@ -0,0 +1,34 @@
{
"go.lintFlags": [
"-checks=\"all,-ST1005\""
],
"go.testTimeout": "300s",
"go.testEnvVars": {
"INSTALLER_LOG_FILE": "${workspaceFolder}/build/log.json"
},
"go.testExplorer.enable": true,
"search.useIgnoreFiles": true,
"search.exclude": {
"**/frontend/coverage/**": true,
"**/frontend/dist/**": true,
"**/frontend/node_modules/**": true,
},
"editor.codeActionsOnSave": {
"source.fixAll.eslint": "explicit"
},
"eslint.validate": [
"javascript",
"javascriptreact",
"typescript",
"typescriptreact"
],
"editor.formatOnSave": false,
"[typescript]": {
"editor.defaultFormatter": "esbenp.prettier-vscode",
"editor.formatOnSave": true
},
"[typescriptreact]": {
"editor.defaultFormatter": "esbenp.prettier-vscode",
"editor.formatOnSave": true
},
}
CLAUDE.md
+150
View File
@@ -0,0 +1,150 @@
# CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
## Core Interaction Rules
1. **Always use English** for all interactions, responses, explanations, and questions with users.
2. **Password Complexity Requirements**: For all password-related development (registration, password reset, API token generation, etc.), the following rules must be enforced:
- Minimum 12 characters
- Must contain at least 1 uppercase letter, 1 lowercase letter, 1 number, and 1 special character
- Common weak passwords (e.g., `password`, `123456`) are prohibited
- Both backend and frontend validation must be implemented; do not rely on frontend validation alone
## Project Overview
**PentAGI** is an automated security testing platform powered by AI agents. It runs autonomous penetration testing workflows using a multi-agent system (Researcher, Developer, Executor agents) that coordinates LLM providers, Docker-sandboxed tool execution, and a persistent vector memory store.
The application is a monorepo with:
- **`backend/`** — Go REST + GraphQL API server
- **`frontend/`** — React + TypeScript web UI
- **`observability/`** — Optional monitoring stack configs
## Build & Development Commands
### Backend (run from `backend/`)
```bash
go mod download # Install dependencies
go build -trimpath -o pentagi ./cmd/pentagi # Build main binary
go test ./... # Run all tests
go test ./pkg/foo/... -v -run TestName # Run specific test
golangci-lint run --timeout=5m # Lint
# Code generation (run after schema changes)
go run github.com/99designs/gqlgen --config ./gqlgen/gqlgen.yml # GraphQL resolvers
swag init -g ../../pkg/server/router.go -o pkg/server/docs/ --parseDependency --parseInternal --parseDepth 2 -d cmd/pentagi # Swagger docs
```
### Frontend (run from `frontend/`)
```bash
pnpm install # Install dependencies
pnpm run dev # Dev server on http://localhost:8000
pnpm run build # Production build
pnpm run lint # ESLint check
pnpm run lint:fix # ESLint auto-fix
pnpm run prettier # Prettier check
pnpm run prettier:fix # Prettier auto-format
pnpm run test # Vitest
pnpm run test:coverage # Coverage report
pnpm run graphql:generate # Regenerate GraphQL types from schema
```
### Docker (run from repo root)
```bash
docker compose up -d # Start core services
docker compose -f docker-compose.yml -f docker-compose-observability.yml up -d # + monitoring
docker compose -f docker-compose.yml -f docker-compose-langfuse.yml up -d # + LLM analytics
docker compose -f docker-compose.yml -f docker-compose-graphiti.yml up -d # + knowledge graph
docker build -t local/pentagi:latest . # Build image
```
The full stack runs at `https://localhost:8443` when using Docker Compose. Copy `.env.example` to `.env` and fill in at minimum the database and at least one LLM provider key.
## Architecture
### Backend Package Structure
| Package | Role |
|---|---|
| `cmd/pentagi/` | Main entry point; initializes config, DB, server |
| `pkg/config/` | Environment-based config parsing |
| `pkg/server/` | Gin router, middleware, auth (JWT/OAuth2/API tokens), Swagger |
| `pkg/controller/` | Business logic for REST endpoints |
| `pkg/graph/` | gqlgen GraphQL schema (`schema.graphqls`) and resolvers |
| `pkg/database/` | GORM models, SQLC queries, goose migrations |
| `pkg/providers/` | LLM provider adapters (OpenAI, Anthropic, Gemini, Bedrock, Ollama, etc.) |
| `pkg/tools/` | Penetration testing tool integrations |
| `pkg/docker/` | Docker SDK wrapper for sandboxed container execution |
| `pkg/terminal/` | Terminal session and command execution management |
| `pkg/queue/` | Async task queue |
| `pkg/csum/` | Chain summarization for LLM context management |
| `pkg/graphiti/` | Knowledge graph (Neo4j via Graphiti) integration |
| `pkg/observability/` | OpenTelemetry tracing, metrics, structured logging |
Database migrations live in `backend/migrations/sql/` and run automatically via goose at startup.
### Frontend Structure
```
frontend/src/
├── app.tsx / main.tsx # Entry points and router setup
├── pages/ # Route-level page components
│ ├── flows/ # Flow management UI
│ └── settings/ # Provider, prompt, token settings
├── components/
│ ├── layouts/ # App shell layouts
│ └── ui/ # Base Radix UI components
├── graphql/ # Auto-generated Apollo types (do not edit)
├── hooks/ # Custom React hooks
├── lib/ # Apollo client, HTTP utilities
└── schemas/ # Zod validation schemas
```
State is managed primarily through Apollo Client (GraphQL) with real-time updates via GraphQL subscriptions over WebSocket.
### Data Flow
1. User creates a "flow" (penetration test) via the UI or REST API.
2. The backend queues the flow and spawns agent goroutines.
3. The Researcher agent gathers information; the Developer plans attack strategies; the Executor runs tools in isolated Docker containers.
4. Results, tool outputs, and LLM reasoning are stored in PostgreSQL (with pgvector for semantic search/memory).
5. Real-time progress is pushed to the frontend via GraphQL subscriptions.
### Authentication
- **Session cookies** for browser login (secure, httpOnly)
- **OAuth2** via Google and GitHub
- **Bearer tokens** (API tokens table) for programmatic API access
### Key Integrations
- **LLM Providers**: OpenAI, Anthropic, Gemini, AWS Bedrock, Ollama, DeepSeek, GLM, Kimi, Qwen, and custom HTTP endpoints — configured via environment variables or the Settings UI
- **Search**: DuckDuckGo, Google, Tavily, Traversaal, Perplexity, Searxng
- **Databases**: PostgreSQL + pgvector (required), Neo4j (optional, for knowledge graph)
- **Observability**: OpenTelemetry → VictoriaMetrics + Loki + Jaeger → Grafana; Langfuse for LLM analytics
### Adding a New LLM Provider
1. Create `backend/pkg/providers/<name>/<name>.go` implementing the `provider.Provider` interface.
2. Add a new `Provider<Name> ProviderType` constant and `DefaultProviderName<Name>` in `pkg/providers/provider/provider.go`.
3. Register the provider in `pkg/providers/providers.go` (`DefaultProviderConfig`, `NewProvider`, `buildProviderFromConfig`, `GetProvider`).
4. Add the new type to the `Valid()` whitelist in `pkg/server/models/providers.go` — **without this step, the REST API returns 422 Unprocessable Entity**.
5. Add the env var key to `pkg/config/config.go` (e.g., `<NAME>_API_KEY`, `<NAME>_SERVER_URL`).
6. Add the new `PROVIDER_TYPE` enum value via a goose migration in `backend/migrations/sql/`.
7. Add the provider icon in `frontend/src/components/icons/<name>.tsx` and register it in `frontend/src/components/icons/provider-icon.tsx`.
8. Update the GraphQL schema/types and frontend settings page if needed.
### Code Generation
When modifying `backend/pkg/graph/schema.graphqls`, re-run the gqlgen command to regenerate resolver stubs. When modifying REST handler annotations, re-run swag to update Swagger docs. When modifying `frontend/src/graphql/*.graphql` query files, re-run `pnpm run graphql:generate` to update TypeScript types.
### Utility Binaries
The backend contains helper binaries for development/testing:
- `cmd/ctester/` — tests container execution
- `cmd/ftester/` — tests LLM function/tool calling
- `cmd/etester/` — tests embedding providers
- `cmd/installer/` — interactive TUI wizard for guided deployment setup (configures `.env`, Docker Compose, DB, search engines, etc.)
CONTRIBUTING.md
+67
View File
@@ -0,0 +1,67 @@
# License Compliance Guide
## Overview
PentAGI is **MIT licensed** and all dependencies use MIT-compatible licenses.
## For Developers
### Adding New Dependencies
When adding new dependencies, ensure they use compatible licenses:
#### Approved Licenses
- MIT
- Apache-2.0
- BSD-2-Clause, BSD-3-Clause
- ISC
- MPL-2.0 (if used without modification)
- 0BSD (public domain)
#### Incompatible Licenses
- GPL, LGPL, AGPL (without special exception)
- CC-BY-SA (for code, OK for data)
- Proprietary/Commercial licenses
### Before Merging PR
1. Update dependencies:
```bash
cd backend && go mod tidy
cd ../frontend && pnpm install
```
2. Generate license reports:
```bash
./scripts/generate-licenses.sh
```
This script automatically collects license information from all dependencies and saves them to the `licenses/` directory. See [licenses/README.md](../licenses/README.md) for details.
3. Scan for issues:
```bash
osv-scanner scan --experimental-licenses="MIT,Apache-2.0,BSD-2-Clause,BSD-3-Clause,ISC,MPL-2.0" backend
osv-scanner scan --experimental-licenses="MIT,Apache-2.0,BSD-2-Clause,BSD-3-Clause,ISC,MPL-2.0" frontend
```
### License Verification Tools
- **osv-scanner** - Security and license scanning (recommended)
- **license-checker** - npm license verification
- **go-licenses** - Go module license extraction (install: `go install github.com/google/go-licenses@latest`)
- **go list** - Go module inspection
## For Docker Builds
License reports are automatically generated during Docker builds:
- Backend reports in `/opt/pentagi/licenses/backend/`
- `dependencies.txt` - List of all Go modules
- `licenses.csv` - Detailed license information (generated by go-licenses)
- Frontend reports in `/opt/pentagi/licenses/frontend/`
- `dependencies.json` - Complete npm dependency tree
- `licenses.json` - Detailed license data
- `licenses.csv` - License summary
## Questions?
Contact: **info@pentagi.com** or **info@vxcontrol.com**
CONTRIBUTORS.md
+138
View File
@@ -0,0 +1,138 @@
# Contributors
This document recognizes all individuals who have contributed to PentAGI. Their work across 370+ commits over 18 months has shaped this project into what it is today.
## Core Team
### Project Lead & Backend Development
- [@asdek](https://github.com/asdek) (Dmitry Nagibin) - Architecture, backend infrastructure, agent system, provider integrations, observability, and project coordination
### Frontend Development
- [@sirozha](https://github.com/sirozha) (Sergey Kozyrenko) - React UI implementation, settings interfaces, GraphQL integration, and frontend architecture
### Backend Development
- [@zavgorodnii](https://github.com/zavgorodnii) (Andrei Zavgorodnii) - Graphiti integration, patch refiner, knowledge graph implementation
---
## External Contributors
We are deeply grateful to the following individuals for their contributions to PentAGI:
### Feature Contributors
#### [@mason5052](https://github.com/mason5052) (Mason Kim)
**Contributions:** Significant testing infrastructure improvements and bug fixes
- Added comprehensive unit test coverage across multiple packages ([PR#199](https://github.com/vxcontrol/pentagi/pull/199), [PR#198](https://github.com/vxcontrol/pentagi/pull/198), [PR#200](https://github.com/vxcontrol/pentagi/pull/200), [PR#201](https://github.com/vxcontrol/pentagi/pull/201), [PR#202](https://github.com/vxcontrol/pentagi/pull/202), [PR#214](https://github.com/vxcontrol/pentagi/pull/214), [PR#213](https://github.com/vxcontrol/pentagi/pull/213))
- Test coverage for: config, version, terminal, server response, embeddings, graph context, executor helpers, custom JSON types ([PR#170](https://github.com/vxcontrol/pentagi/pull/170)), context registry ([PR#171](https://github.com/vxcontrol/pentagi/pull/171)), executor terminal utilities ([PR#172](https://github.com/vxcontrol/pentagi/pull/172)), search tools ([PR#153](https://github.com/vxcontrol/pentagi/pull/153))
- Test coverage for LLM providers (DeepSeek, GLM, Kimi, Qwen) ([PR#189](https://github.com/vxcontrol/pentagi/pull/189))
- Fixed critical bugs:
- Detached command context isolation ([PR#179](https://github.com/vxcontrol/pentagi/pull/179))
- Agent chain iteration cap and repeating detector ([PR#178](https://github.com/vxcontrol/pentagi/pull/178), [PR#180](https://github.com/vxcontrol/pentagi/pull/180))
- HTTP client global mutation in search tools ([PR#151](https://github.com/vxcontrol/pentagi/pull/151))
- Browser graceful screenshot failure handling ([PR#150](https://github.com/vxcontrol/pentagi/pull/150))
- Silent error handling in GetTool ([PR#152](https://github.com/vxcontrol/pentagi/pull/152))
- OAuth callback missing return ([PR#127](https://github.com/vxcontrol/pentagi/pull/127))
- OAuth state field validation ([PR#125](https://github.com/vxcontrol/pentagi/pull/125), [PR#120](https://github.com/vxcontrol/pentagi/pull/120))
- Langfuse TLS configuration respect ([PR#132](https://github.com/vxcontrol/pentagi/pull/132))
- CA private key cleanup in entrypoint ([PR#168](https://github.com/vxcontrol/pentagi/pull/168))
- Google search options unused ([PR#167](https://github.com/vxcontrol/pentagi/pull/167))
- Resource leaks and unbounded allocation in tools ([PR#141](https://github.com/vxcontrol/pentagi/pull/141))
- Debug console.log removal ([PR#141](https://github.com/vxcontrol/pentagi/pull/141))
- Swagger missing quote ([PR#140](https://github.com/vxcontrol/pentagi/pull/140))
- Terminal typo correction ([PR#164](https://github.com/vxcontrol/pentagi/pull/164))
- Documentation improvements: fixed typos and grammar issues ([PR#121](https://github.com/vxcontrol/pentagi/pull/121))
#### [@niuqun2003](https://github.com/niuqun2003) (niuqun2003)
**Contributions:** Chinese LLM provider ecosystem integration ([PR#154](https://github.com/vxcontrol/pentagi/pull/154))
- Added support for DeepSeek, GLM, Kimi, and Qwen LLM providers
- Implemented provider configurations and API integrations
#### [@Priyanka-2725](https://github.com/Priyanka-2725) (Priyanka Singh)
**Contributions:** AWS Bedrock provider improvements, Sploitus inegration implementation and bug fixes
- Fixed Bedrock toolConfig runtime failures ([PR#166](https://github.com/vxcontrol/pentagi/pull/166))
- Implemented better exploit finding capabilities by Sploitus ([PR#133](https://github.com/vxcontrol/pentagi/pull/133))
- Fixed terminal command handling logic ([PR#124](https://github.com/vxcontrol/pentagi/pull/124))
#### [@Alex-wuhu](https://github.com/Alex-wuhu) (Alex)
**Contributions:** Novita AI integration ([PR#162](https://github.com/vxcontrol/pentagi/pull/162))
- Added Novita AI as optional LLM provider
- Updated Novita models configuration
#### [@efe-arv](https://github.com/efe-arv) (Efe Büken)
**Contributions:** HTTP client timeout configuration ([PR#205](https://github.com/vxcontrol/pentagi/pull/205))
- Added configurable timeout to HTTP client for improved reliability
#### [@manusjs](https://github.com/manusjs) (manusjs)
**Contributions:** Bedrock tool configuration fix ([PR#196](https://github.com/vxcontrol/pentagi/pull/196))
- Fixed Bedrock to always include toolConfig when messages contain toolUse/toolResult blocks
#### [@stoykovstoyk](https://github.com/stoykovstoyk) (Stoyko Stoykov)
**Contributions:** SearXNG meta search engine integration ([PR#53](https://github.com/vxcontrol/pentagi/pull/53))
- Added SearXNG as a search engine option
- Implemented tool integration for meta search capabilities
#### [@kaikreuzer](https://github.com/kaikreuzer) (Kai Kreuzer)
**Contributions:** AWS credentials support ([PR#90](https://github.com/vxcontrol/pentagi/pull/90))
- Added support for temporary AWS credentials
- Enhanced AWS session token handling
#### [@mrigankad](https://github.com/mrigankad) (Mriganka Dey)
**Contributions:** Security and bug fixes ([PR#104](https://github.com/vxcontrol/pentagi/pull/104))
- Various security improvements and bug fixes across the codebase
#### [@salmanmkc](https://github.com/salmanmkc) (Salman Chishti)
**Contributions:** GitHub Actions modernization
- Upgraded GitHub Actions to latest versions ([PR#112](https://github.com/vxcontrol/pentagi/pull/112))
- Upgraded GitHub Actions for Node 24 compatibility ([PR#111](https://github.com/vxcontrol/pentagi/pull/111))
### Bug Fixes & Improvements
#### [@Vaibhavee89](https://github.com/Vaibhavee89) (Vaibhavee Singh)
**Contributions:** Documentation enhancement
- Added external network access configuration guide to README
#### [@PeterDaveHello](https://github.com/PeterDaveHello) (Peter Dave Hello)
**Contributions:** Dockerfile optimization ([PR#50](https://github.com/vxcontrol/pentagi/pull/50))
- Removed unnecessary cleanup steps in Dockerfile for improved build efficiency
#### [@s-b-repo](https://github.com/s-b-repo) (S.B)
**Contributions:** Security improvements
- Implemented file size limit and path escaping for enhanced security
- Fixed typo in executor.go ('Incoming')
#### [@SkyFlyingMouse](https://github.com/SkyFlyingMouse) (SkyFlyingMouse)
**Contributions:** Code quality ([PR#128](https://github.com/vxcontrol/pentagi/pull/128))
- Fixed Docker client constant name typo in backend
#### [@haosenwang1018](https://github.com/haosenwang1018) (Sense_wang)
**Contributions:** Development environment ([PR#163](https://github.com/vxcontrol/pentagi/pull/163))
- Expanded .gitignore with IDE and OS patterns for cleaner repository
#### [@hhktony](https://github.com/hhktony) (Tony Xu)
**Contributions:** Documentation ([PR#32](https://github.com/vxcontrol/pentagi/pull/32))
- README.md improvements and clarifications
---
## Acknowledgments
This project exists because of the collective effort of everyone listed above. From major feature implementations to small bug fixes, every contribution has made PentAGI better.
---
## How to Contribute
Interested in contributing to PentAGI? We welcome contributions of all kinds:
- Bug reports and fixes
- New features and enhancements
- Documentation improvements
- Testing and QA
- Code reviews
Please see our [CONTRIBUTING.md](CONTRIBUTING.md) guide for more information.
---
**Note:** Due to repository history rewriting on March 29, 2026, to resolve licensing matters, individual commit history is no longer visible in GitHub's interface. This document preserves the record of all contributions made during the project's development from January 2025 to March 2026.
Dockerfile
+199
View File
@@ -0,0 +1,199 @@
# syntax=docker/dockerfile:1.4
# ========================================
# Stage 1: Frontend Application Build
# ========================================
FROM node:23-slim AS frontend-compiler
# Production build configuration
ENV NODE_ENV=production
ENV VITE_BUILD_MEMORY_LIMIT=4096
ENV NODE_OPTIONS="--max-old-space-size=4096"
ENV PNPM_HOME="/usr/local/share/pnpm"
ENV PATH="$PNPM_HOME:$PATH"
WORKDIR /app/ui
# Install build essentials and enable pnpm via corepack
RUN apt-get update && apt-get install -y \
ca-certificates \
tzdata \
gcc \
g++ \
make \
git \
&& corepack enable && corepack prepare pnpm@latest --activate
# GraphQL schema for code generation
COPY ./backend/pkg/graph/schema.graphqls ../backend/pkg/graph/
# Application source code
COPY frontend/ .
# Install dependencies
RUN --mount=type=cache,target=/root/.local/share/pnpm/store \
pnpm install --frozen-lockfile
# Generate license report for frontend dependencies
RUN pnpm add -g license-checker && \
mkdir -p /licenses/frontend && \
license-checker --production --json > /licenses/frontend/licenses.json && \
license-checker --production --csv > /licenses/frontend/licenses.csv
# Build frontend with optimizations and parallel processing
RUN pnpm run build -- \
--mode production \
--minify esbuild \
--outDir dist \
--emptyOutDir \
--sourcemap false \
--target es2020
# ========================================
# Stage 2: Backend Services Compilation
# ========================================
FROM golang:1.24-bookworm AS api-builder
# Version injection arguments
ARG PACKAGE_VER=develop
ARG PACKAGE_REV=
# Static binary compilation settings
ENV CGO_ENABLED=0
ENV GO111MODULE=on
# Install compilation toolchain and dependencies
RUN apt-get update && apt-get install -y \
ca-certificates \
tzdata \
gcc \
g++ \
make \
git \
musl-dev
WORKDIR /app/backend
COPY backend/ .
# Fetch Go module dependencies (cached for faster rebuilds)
RUN --mount=type=cache,target=/go/pkg/mod \
go mod download && go mod verify
# Install go-licenses tool for license extraction
RUN --mount=type=cache,target=/go/pkg/mod \
go install github.com/google/go-licenses@latest
# Generate license reports for backend dependencies
RUN mkdir -p /licenses/backend && \
go list -m all > /licenses/backend/dependencies.txt && \
GOROOT=$(go env GOROOT) GOTOOLCHAIN=auto go-licenses csv ./cmd/pentagi > /licenses/backend/licenses.csv 2>/dev/null || true
# Compile main application binary with embedded version metadata
RUN go build -trimpath \
-ldflags "\
-X pentagi/pkg/version.PackageName=pentagi \
-X pentagi/pkg/version.PackageVer=${PACKAGE_VER} \
-X pentagi/pkg/version.PackageRev=${PACKAGE_REV}" \
-o /pentagi ./cmd/pentagi
# Build ctester utility
RUN go build -trimpath \
-ldflags "\
-X pentagi/pkg/version.PackageName=ctester \
-X pentagi/pkg/version.PackageVer=${PACKAGE_VER} \
-X pentagi/pkg/version.PackageRev=${PACKAGE_REV}" \
-o /ctester ./cmd/ctester
# Build ftester utility
RUN go build -trimpath \
-ldflags "\
-X pentagi/pkg/version.PackageName=ftester \
-X pentagi/pkg/version.PackageVer=${PACKAGE_VER} \
-X pentagi/pkg/version.PackageRev=${PACKAGE_REV}" \
-o /ftester ./cmd/ftester
# Build etester utility
RUN go build -trimpath \
-ldflags "\
-X pentagi/pkg/version.PackageName=etester \
-X pentagi/pkg/version.PackageVer=${PACKAGE_VER} \
-X pentagi/pkg/version.PackageRev=${PACKAGE_REV}" \
-o /etester ./cmd/etester
# ========================================
# Stage 3: Production Runtime Environment
# ========================================
FROM alpine:3.23.3
# Establish non-privileged execution context with docker socket access
RUN addgroup -g 998 docker && \
addgroup -S pentagi && \
adduser -S pentagi -G pentagi && \
addgroup pentagi docker
# Install required packages
RUN apk --no-cache add ca-certificates openssl openssh-keygen shadow
ADD scripts/entrypoint.sh /opt/pentagi/bin/
RUN sed -i 's/\r//' /opt/pentagi/bin/entrypoint.sh && \
chmod +x /opt/pentagi/bin/entrypoint.sh
RUN mkdir -p \
/root/.ollama \
/opt/pentagi/bin \
/opt/pentagi/ssl \
/opt/pentagi/fe \
/opt/pentagi/logs \
/opt/pentagi/data \
/opt/pentagi/conf && \
chmod 777 /root/.ollama
COPY --from=api-builder /pentagi /opt/pentagi/bin/pentagi
COPY --from=api-builder /ctester /opt/pentagi/bin/ctester
COPY --from=api-builder /ftester /opt/pentagi/bin/ftester
COPY --from=api-builder /etester /opt/pentagi/bin/etester
COPY --from=frontend-compiler /app/ui/dist /opt/pentagi/fe
COPY --from=api-builder /licenses/backend /opt/pentagi/licenses/backend
COPY --from=frontend-compiler /licenses/frontend /opt/pentagi/licenses/frontend
# Copy provider configuration files
COPY examples/configs/azure-openai.provider.yml /opt/pentagi/conf/
COPY examples/configs/custom-openai.provider.yml /opt/pentagi/conf/
COPY examples/configs/deepinfra.provider.yml /opt/pentagi/conf/
COPY examples/configs/deepseek.provider.yml /opt/pentagi/conf/
COPY examples/configs/moonshot.provider.yml /opt/pentagi/conf/
COPY examples/configs/ollama-cloud.provider.yml /opt/pentagi/conf/
COPY examples/configs/ollama-llama318b-instruct.provider.yml /opt/pentagi/conf/
COPY examples/configs/ollama-llama318b.provider.yml /opt/pentagi/conf/
COPY examples/configs/ollama-qwen332b-fp16-tc.provider.yml /opt/pentagi/conf/
COPY examples/configs/ollama-qwq32b-fp16-tc.provider.yml /opt/pentagi/conf/
COPY examples/configs/openrouter.provider.yml /opt/pentagi/conf/
COPY examples/configs/novita.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen3.5-27b-fp8-no-think.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen3.5-27b-fp8.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen3.6-27b-fp8-no-think.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen3.6-27b-fp8.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen3.6-35b-a3b-fp8-no-think.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen3.6-35b-a3b-fp8.provider.yml /opt/pentagi/conf/
COPY examples/configs/vllm-qwen332b-fp16.provider.yml /opt/pentagi/conf/
COPY LICENSE /opt/pentagi/LICENSE
COPY NOTICE /opt/pentagi/NOTICE
COPY EULA.md /opt/pentagi/EULA
COPY EULA.md /opt/pentagi/fe/EULA.md
RUN chown -R pentagi:pentagi /opt/pentagi
WORKDIR /opt/pentagi
USER pentagi
ENTRYPOINT ["/opt/pentagi/bin/entrypoint.sh", "/opt/pentagi/bin/pentagi"]
# Image Metadata
LABEL org.opencontainers.image.source="https://github.com/vxcontrol/pentagi"
LABEL org.opencontainers.image.description="Fully autonomous AI Agents system capable of performing complex penetration testing tasks"
LABEL org.opencontainers.image.authors="PentAGI Development Team"
LABEL org.opencontainers.image.licenses="MIT License"
EULA.md
+93
View File
@@ -0,0 +1,93 @@
# PentAGI End User License Agreement
## Introduction
This **End User License Agreement (EULA)** governs the terms and conditions for the use of PentAGI, an advanced AI-powered penetration testing tool. This product is provided by the **PentAGI Development Team**, and is distributed in the form of [source code](https://github.com/vxcontrol/pentagi) available on GitHub under the MIT license as well as [pre-built Docker images](https://hub.docker.com/r/vxcontrol/pentagi) available on Docker Hub.
Users agree to this EULA when downloading either the source code or the Docker images or by accessing the product's interface through its web UI. It is the user's responsibility to ensure compliance with all applicable laws and standards when utilizing PentAGI. This product is intended for lawful penetration testing purposes and research purposes only and does not inherently possess tools used for executing cyber attacks. Instead, it facilitates the download of publicly available penetration testing tools such as those from Kali Linux or other similar distributions.
PentAGI operates independently of services provided by the Developers and allows users to self-deploy all components. Users initiate interaction through a web user interface, which is part of the product itself. Integration with external LLM providers and search systems requires careful oversight by the user to ensure data compliance, including regulations like GDPR.
The **PentAGI Development Team** can be contacted via GitHub or through the email address [info@pentagi.com](mailto:info@pentagi.com). This document should be reviewed in its entirety to fully understand the terms and legal obligations therein.
## License Grant
Under this EULA, the **PentAGI Development Team** grants you a non-exclusive, non-transferable, revocable license to use the PentAGI software solely for lawful penetration testing purposes. This license is effective when you download the source code or Docker images and remains in effect until terminated as outlined in this agreement.
The source code of PentAGI is provided under the MIT license, the terms of which are incorporated herein by reference. This EULA governs your use of the PentAGI software as a whole, including any pre-built Docker images and the web UI, and applies in addition to the MIT license. In the event of any conflict between this EULA and the MIT license, the terms of the MIT license shall prevail with respect to the source code.
You are permitted to use the PentAGI software on your own infrastructure, self-deploying all components according to provided documentation. The license covers usage as allowed by the MIT license under which the source code is distributed, but does not extend to any proprietary tools that may be downloaded or used in conjunction with the PentAGI software.
You may not sublicense, sell, lease, or distribute the PentAGI software or its derivatives in any form other than stated in the license agreement. Modification and redistribution are permitted under the MIT license conditions; however, the **PentAGI Development Team** holds no responsibility for any alterations not published by them through the official GitHub or Docker Hub pages.
## Acceptable Use
PentAGI is to be used exclusively for authorized penetration testing and security assessments in environments where you have explicit permission from the network owner. You must ensure that all usage complies with applicable laws, standards, and regulations, particularly those concerning cybersecurity and data protection.
You are solely responsible for the execution and outcomes of any tasks set for AI agents within the PentAGI interface. The logic and actions of the AI agents are strictly determined by the tasks and instructions you provide. The **PentAGI Development Team** does not supervise or control the actions of the AI agents and is not responsible for any consequences arising from their actions. You must verify that all data sent to AI agents, external LLM providers, search systems, or stored within PentAGI complies with legal standards and regulations, including but not limited to GDPR.
You must not use PentAGI in any critical infrastructure, emergency response systems, or other high-risk environments without proper testing and validation. The software is intended for research and testing purposes only and should not be deployed in production environments without thorough security assessment.
Using PentAGI for any activity that violates laws or regulations, including but not limited to unauthorized network access, is strictly prohibited. Users found using the software for illegal purposes may have their license revoked and could face further legal consequences, as determined by law enforcement.
## Data Privacy and Security
You acknowledge that PentAGI may process sensitive information during penetration testing activities. You are solely responsible for ensuring that all data processing complies with applicable privacy laws and regulations, including GDPR, CCPA, and other relevant data protection regulations.
The **PentAGI Development Team** does not collect, store, or process any user data through the software. All data processing occurs locally within your infrastructure or through third-party services that you configure. You are responsible for implementing appropriate security measures to protect any sensitive data processed through PentAGI.
When using PentAGI's integration capabilities with external services, you must ensure that all data transfers comply with applicable data protection regulations and that you have obtained necessary consents for data processing.
## Third-Party Services
PentAGI integrates with external third-party services, including but not limited to Large Language Model (LLM) providers such as OpenAI, Anthropic, Deep Infra, OpenRouter, and search engines such as Tavily, Traversaal, Perplexity, DuckDuckGo, Google, Sploitus and Searxng. You acknowledge and agree that your use of these third-party services is at your sole discretion and responsibility.
When using self-hosted or local LLM servers compatible with OpenAI API, you are solely responsible for ensuring the security and compliance of these deployments. The PentAGI Development Team bears no responsibility for any data leaks or security issues arising from the use of such local deployments.
The **PentAGI Development Team** does not control and is not responsible for any content, data, or privacy practices of these third-party services. You are responsible for ensuring that your use of these services, including any data you transmit to them, complies with all applicable laws and regulations, including data protection and privacy laws such as the General Data Protection Regulation (GDPR).
By using PentAGI's integration with third-party services, you agree to comply with any terms and conditions imposed by those services. The **PentAGI Development Team** disclaims any and all liability arising from your use of third-party services and makes no representations or warranties regarding the functionality or security of these services.
## Disclaimer of Warranties
PentAGI is provided "as is" and "as available," with all faults and without warranty of any kind. To the maximum extent permitted by applicable law, the **PentAGI Development Team** disclaims all warranties, whether express, implied, statutory, or otherwise, regarding the software, including without limitation any warranties of merchantability, fitness for a particular purpose, title, and non-infringement.
The **PentAGI Development Team** disclaims any liability for actions performed by AI agents within the software, or for any data transmitted to third-party services by the user.
The Developers do not warrant that the PentAGI software will operate uninterrupted or error-free, that defects will be corrected, or that the software is free of viruses or other harmful components. Your use of the software is at your sole risk, and you assume full responsibility for any costs or losses incurred.
## Limitation of Liability
To the fullest extent permitted by law, in no event shall the **PentAGI Development Team** be liable for any direct, indirect, incidental, special, consequential, or punitive damages, including but not limited to lost profits, lost savings, business interruption, or loss of data, arising out of your use or inability to use the PentAGI software, even if advised of the possibility of such damages.
The **PentAGI Development Team** shall not be liable for any damages or losses resulting from the actions of AI agents operated through PentAGI, or from the use of third-party services integrated with PentAGI.
The **PentAGI Development Team** shall not be liable for any damages or losses resulting from modifications to the source code, whether made by you or third parties, including but not limited to forks of the GitHub repository or modified Docker images not officially published by the PentAGI Development Team.
The total cumulative liability of the **PentAGI Development Team** arising from or related to this EULA, whether in contract, tort, or otherwise, shall not exceed the amount paid by you for the software.
## Indemnification
You agree to indemnify, defend, and hold harmless the **PentAGI Development Team**, its members, and any of its contractors, suppliers, or affiliates from and against any and all claims, liabilities, damages, losses, or expenses, including reasonable attorneys' fees and costs, arising out of or in any way connected to your use of the PentAGI software, your violation of this EULA, or your violation of any law or the rights of a third party.
## Termination
This EULA is effective until terminated either by you or by the **PentAGI Development Team**. You may terminate this agreement at any time by ceasing all use of the PentAGI software and destroying all copies in your possession.
The **PentAGI Development Team** reserves the right to terminate this EULA and your access to the software immediately, without notice, if you breach any term of this agreement. Upon termination, you must cease all use of the software and destroy all copies, whether full or partial, in your possession.
## Governing Law and Dispute Resolution
This EULA and any disputes arising out of or related to it shall be governed by and construed in accordance with the laws of the United Kingdom, without regard to its conflict of law principles.
Any and all disputes arising under or in connection with this EULA shall be resolved through negotiations. If the parties cannot resolve a dispute through good-faith negotiations within 90 days, they agree to submit the dispute to binding arbitration under the rules of an arbitration body in the United Kingdom. The language of arbitration shall be English.
## Miscellaneous Provisions
This EULA constitutes the entire agreement between you and the **PentAGI Development Team** regarding the use of PentAGI and supersedes all prior agreements and understandings. If any provision of this EULA is found to be invalid or unenforceable, the remainder shall continue to be fully enforceable and effective.
The **PentAGI Development Team** publishes official updates and versions of the software only on the GitHub repository at [vxcontrol/pentagi](https://github.com/vxcontrol/pentagi) and on Docker Hub at [vxcontrol/pentagi](https://hub.docker.com/r/vxcontrol/pentagi). Any forks, derivative works, or modified versions of the software are not endorsed by the **PentAGI Development Team**, and the team bears no responsibility for such versions.
The Developers reserve the right to modify this EULA at any time by posting the revised EULA on the official PentAGI GitHub page or notifying users via email. Any modifications will be effective immediately upon posting or notification for the next product versions.
Failure by either party to enforce any provision of this EULA shall not constitute a waiver of future enforcement of that or any other provision.
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2025 PentAGI Development Team
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
NOTICE
+9
View File
@@ -0,0 +1,9 @@
PentAGI, Fully autonomous AI Agent capable of performing complex penetration testing tasks.
Copyright 2025 PentAGI Development Team
Licensed under MIT License. See LICENSE and EULA for terms.
NOTICE: This software integrates VXControl Cloud SDK for enhanced intelligence services.
VXControl Cloud SDK is licensed under AGPL-3.0 with a special exception for this official PentAGI project.
For more details, see the License section in README.md and VXControl Cloud SDK license terms.
README.md
+3734
View File
File diff suppressed because it is too large Load Diff
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`vxcontrol/pentagi`
- 原始仓库：https://github.com/vxcontrol/pentagi
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
backend/cmd/ctester/main.go
+368
View File
@@ -0,0 +1,368 @@
package main
import (
"context"
"flag"
"fmt"
"log"
"os"
"strings"
"time"
"pentagi/pkg/config"
"pentagi/pkg/providers/anthropic"
"pentagi/pkg/providers/bedrock"
"pentagi/pkg/providers/custom"
"pentagi/pkg/providers/deepseek"
"pentagi/pkg/providers/gemini"
"pentagi/pkg/providers/glm"
"pentagi/pkg/providers/kimi"
"pentagi/pkg/providers/ollama"
"pentagi/pkg/providers/openai"
"pentagi/pkg/providers/pconfig"
"pentagi/pkg/providers/provider"
"pentagi/pkg/providers/qwen"
"pentagi/pkg/providers/tester"
"pentagi/pkg/providers/tester/testdata"
"pentagi/pkg/version"
"github.com/joho/godotenv"
"github.com/sirupsen/logrus"
)
func main() {
envFile := flag.String("env", ".env", "Path to environment file")
providerType := flag.String("type", "custom", "Provider type [custom, openai, anthropic, gemini, bedrock, ollama, deepseek, glm, kimi, qwen]")
providerName := flag.String("name", "", "Provider name using as PROVDER_NAME/MODEL_NAME while building provider config")
configPath := flag.String("config", "", "Path to provider config file")
testsPath := flag.String("tests", "", "Path to custom tests YAML file")
reportPath := flag.String("report", "", "Path to write report file")
agentTypes := flag.String("agents", "all", "Comma-separated agent types to test")
testGroups := flag.String("groups", "all", "Comma-separated test groups to run")
workers := flag.Int("workers", 4, "Number of workers to use")
verbose := flag.Bool("verbose", false, "Enable verbose output")
flag.Parse()
logrus.Infof("Starting PentAGI Provider Configuration Tester %s", version.GetBinaryVersion())
if err := godotenv.Load(*envFile); err != nil {
log.Println("Warning: Error loading .env file:", err)
}
cfg, err := config.NewConfig()
if err != nil {
log.Fatalf("Error loading config: %v", err)
}
if *configPath != "" {
cfg.LLMServerConfig = *configPath
cfg.OllamaServerConfig = *configPath
}
if *providerName != "" {
cfg.LLMServerProvider = *providerName
}
prv, err := createProvider(*providerType, cfg)
if err != nil {
log.Fatalf("Error creating provider: %v", err)
}
fmt.Printf("Testing %s Provider\n", *providerType)
fmt.Println("=================================================")
var testOptions []tester.TestOption
if *agentTypes != "all" {
selectedTypes := parseAgentTypes(strings.Split(*agentTypes, ","))
testOptions = append(testOptions, tester.WithAgentTypes(selectedTypes...))
}
if *testGroups != "all" {
selectedGroups := parseTestGroups(strings.Split(*testGroups, ","))
testOptions = append(testOptions, tester.WithGroups(selectedGroups...))
} else {
// Include all available groups when "all" is specified
allGroups := []testdata.TestGroup{
testdata.TestGroupBasic,
testdata.TestGroupAdvanced,
testdata.TestGroupJSON,
testdata.TestGroupKnowledge,
}
testOptions = append(testOptions, tester.WithGroups(allGroups...))
}
if *testsPath != "" {
registry, err := loadCustomTests(*testsPath)
if err != nil {
log.Fatalf("Error loading custom tests: %v", err)
}
testOptions = append(testOptions, tester.WithCustomRegistry(registry))
}
testOptions = append(
testOptions,
tester.WithVerbose(*verbose),
tester.WithParallelWorkers(*workers),
)
results, err := tester.TestProvider(context.Background(), prv, testOptions...)
if err != nil {
log.Fatalf("Error running tests: %v", err)
}
agentResults := convertToAgentResults(results, prv)
PrintSummaryReport(agentResults)
if *reportPath != "" {
if err := WriteReportToFile(agentResults, *reportPath); err != nil {
log.Printf("Error writing report: %v", err)
} else {
fmt.Printf("Report written to %s\n", *reportPath)
}
}
}
func createProvider(providerType string, cfg *config.Config) (provider.Provider, error) {
switch providerType {
case "custom":
providerConfig, err := custom.DefaultProviderConfig(cfg)
if err != nil {
return nil, fmt.Errorf("error creating custom provider config: %w", err)
}
return custom.New(cfg, provider.DefaultProviderNameCustom, providerConfig)
case "openai":
if cfg.OpenAIKey == "" {
return nil, fmt.Errorf("OpenAI key is not set")
}
providerConfig, err := openai.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating openai provider config: %w", err)
}
return openai.New(cfg, provider.DefaultProviderNameOpenAI, providerConfig)
case "anthropic":
if cfg.AnthropicAPIKey == "" {
return nil, fmt.Errorf("Anthropic API key is not set")
}
providerConfig, err := anthropic.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating anthropic provider config: %w", err)
}
return anthropic.New(cfg, provider.DefaultProviderNameAnthropic, providerConfig)
case "gemini":
if cfg.GeminiAPIKey == "" {
return nil, fmt.Errorf("Gemini API key is not set")
}
providerConfig, err := gemini.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating gemini provider config: %w", err)
}
return gemini.New(cfg, provider.DefaultProviderNameGemini, providerConfig)
case "bedrock":
if !cfg.BedrockDefaultAuth && cfg.BedrockBearerToken == "" &&
(cfg.BedrockAccessKey == "" || cfg.BedrockSecretKey == "") {
return nil, fmt.Errorf("Bedrock requires authentication: set " +
"BEDROCK_DEFAULT_AUTH=true, BEDROCK_BEARER_TOKEN, or " +
"BEDROCK_ACCESS_KEY_ID+BEDROCK_SECRET_ACCESS_KEY")
}
providerConfig, err := bedrock.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating bedrock provider config: %w", err)
}
return bedrock.New(cfg, provider.DefaultProviderNameBedrock, providerConfig)
case "ollama":
if cfg.OllamaServerURL == "" {
return nil, fmt.Errorf("Ollama server URL is not set")
}
providerConfig, err := ollama.DefaultProviderConfig(cfg)
if err != nil {
return nil, fmt.Errorf("error creating ollama provider config: %w", err)
}
return ollama.New(cfg, provider.DefaultProviderNameOllama, providerConfig)
case "deepseek":
if cfg.DeepSeekAPIKey == "" {
return nil, fmt.Errorf("DeepSeek API key is not set")
}
providerConfig, err := deepseek.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating deepseek provider config: %w", err)
}
return deepseek.New(cfg, provider.DefaultProviderNameDeepSeek, providerConfig)
case "glm":
if cfg.GLMAPIKey == "" {
return nil, fmt.Errorf("GLM Zhipu AI API key is not set")
}
providerConfig, err := glm.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating glm provider config: %w", err)
}
return glm.New(cfg, provider.DefaultProviderNameGLM, providerConfig)
case "kimi":
if cfg.KimiAPIKey == "" {
return nil, fmt.Errorf("Kimi Moonshot AI API key is not set")
}
providerConfig, err := kimi.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating kimi provider config: %w", err)
}
return kimi.New(cfg, provider.DefaultProviderNameKimi, providerConfig)
case "qwen":
if cfg.QwenAPIKey == "" {
return nil, fmt.Errorf("Qwen Alibaba Cloud API key is not set")
}
providerConfig, err := qwen.DefaultProviderConfig()
if err != nil {
return nil, fmt.Errorf("error creating qwen provider config: %w", err)
}
return qwen.New(cfg, provider.DefaultProviderNameQwen, providerConfig)
default:
return nil, fmt.Errorf("unsupported provider type: %s", providerType)
}
}
func parseAgentTypes(agentStrings []string) []pconfig.ProviderOptionsType {
var agentTypes []pconfig.ProviderOptionsType
validTypes := map[string]pconfig.ProviderOptionsType{
"simple": pconfig.OptionsTypeSimple,
"simple_json": pconfig.OptionsTypeSimpleJSON,
"primary_agent": pconfig.OptionsTypePrimaryAgent,
"assistant": pconfig.OptionsTypeAssistant,
"generator": pconfig.OptionsTypeGenerator,
"refiner": pconfig.OptionsTypeRefiner,
"adviser": pconfig.OptionsTypeAdviser,
"reflector": pconfig.OptionsTypeReflector,
"searcher": pconfig.OptionsTypeSearcher,
"enricher": pconfig.OptionsTypeEnricher,
"coder": pconfig.OptionsTypeCoder,
"installer": pconfig.OptionsTypeInstaller,
"pentester": pconfig.OptionsTypePentester,
}
for _, agentStr := range agentStrings {
agentStr = strings.TrimSpace(agentStr)
if agentType, ok := validTypes[agentStr]; ok {
agentTypes = append(agentTypes, agentType)
} else {
log.Printf("Warning: Unknown agent type '%s', skipping", agentStr)
}
}
return agentTypes
}
func parseTestGroups(groupStrings []string) []testdata.TestGroup {
var groups []testdata.TestGroup
validGroups := map[string]testdata.TestGroup{
"basic": testdata.TestGroupBasic,
"advanced": testdata.TestGroupAdvanced,
"json": testdata.TestGroupJSON,
"knowledge": testdata.TestGroupKnowledge,
}
for _, groupStr := range groupStrings {
groupStr = strings.TrimSpace(groupStr)
if group, ok := validGroups[groupStr]; ok {
groups = append(groups, group)
} else {
log.Printf("Warning: Unknown test group '%s', skipping", groupStr)
}
}
return groups
}
func convertToAgentResults(results tester.ProviderTestResults, prv provider.Provider) []AgentTestResult {
var agentResults []AgentTestResult
// Create mapping of agent types to their data
agentTypeMap := map[pconfig.ProviderOptionsType]struct {
name string
results tester.AgentTestResults
}{
pconfig.OptionsTypeSimple: {"simple", results.Simple},
pconfig.OptionsTypeSimpleJSON: {"simple_json", results.SimpleJSON},
pconfig.OptionsTypePrimaryAgent: {"primary_agent", results.PrimaryAgent},
pconfig.OptionsTypeAssistant: {"assistant", results.Assistant},
pconfig.OptionsTypeGenerator: {"generator", results.Generator},
pconfig.OptionsTypeRefiner: {"refiner", results.Refiner},
pconfig.OptionsTypeAdviser: {"adviser", results.Adviser},
pconfig.OptionsTypeReflector: {"reflector", results.Reflector},
pconfig.OptionsTypeSearcher: {"searcher", results.Searcher},
pconfig.OptionsTypeEnricher: {"enricher", results.Enricher},
pconfig.OptionsTypeCoder: {"coder", results.Coder},
pconfig.OptionsTypeInstaller: {"installer", results.Installer},
pconfig.OptionsTypePentester: {"pentester", results.Pentester},
}
// Use deterministic order from AllAgentTypes
for _, agentType := range pconfig.AllAgentTypes {
agentData, exists := agentTypeMap[agentType]
if !exists {
continue
}
agentTypeName := agentData.name
agentTestResults := agentData.results
if len(agentTestResults) == 0 {
continue
}
result := AgentTestResult{
AgentType: agentTypeName,
ModelName: prv.Model(agentType),
}
var totalLatency time.Duration
for _, testResult := range agentTestResults {
oldResult := TestResult{
Name: testResult.Name,
Type: string(testResult.Type),
Success: testResult.Success,
Error: testResult.Error,
Streaming: testResult.Streaming,
Reasoning: testResult.Reasoning,
LatencyMs: testResult.Latency.Milliseconds(),
}
if testResult.Group == testdata.TestGroupBasic {
result.BasicTests = append(result.BasicTests, oldResult)
} else {
result.AdvancedTests = append(result.AdvancedTests, oldResult)
}
result.TotalTests++
if testResult.Success {
result.TotalSuccess++
}
if testResult.Reasoning {
result.Reasoning = true
}
totalLatency += testResult.Latency
}
if result.TotalTests > 0 {
result.AverageLatency = totalLatency / time.Duration(result.TotalTests)
}
agentResults = append(agentResults, result)
}
return agentResults
}
func loadCustomTests(path string) (*testdata.TestRegistry, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read tests file: %w", err)
}
return testdata.LoadRegistryFromYAML(data)
}
backend/cmd/ctester/models.go
+30
View File
@@ -0,0 +1,30 @@
package main
import "time"
// TestResult represents the result of a single test for CLI compatibility
type TestResult struct {
Name string
Type string
Success bool
Error error
Streaming bool
Reasoning bool
LatencyMs int64
Response string
Expected string
}
// AgentTestResult collects test results for each agent type for CLI compatibility
type AgentTestResult struct {
AgentType string
ModelName string
Reasoning bool
BasicTests []TestResult
AdvancedTests []TestResult
TotalSuccess int
TotalTests int
AverageLatency time.Duration
SkippedAdvanced bool
SkippedReason string
}
backend/cmd/ctester/report.go
+226
View File
@@ -0,0 +1,226 @@
package main
import (
"fmt"
"os"
"text/tabwriter"
"time"
)
// PrintAgentResults prints the test results for a single agent
func PrintAgentResults(result AgentTestResult) {
fmt.Println("\nTest Results:")
// Basic tests section
if len(result.BasicTests) > 0 {
fmt.Println("\nBasic Tests:")
for _, test := range result.BasicTests {
status := "✓"
if !test.Success {
status = "✗"
}
name := test.Name
if test.Streaming {
name = fmt.Sprintf("Streaming %s", name)
}
fmt.Printf("[%s] %s (%.3fs)\n", status, name, float64(test.LatencyMs)/1000)
if !test.Success && test.Error != nil {
fmt.Printf(" Error: %v\n", test.Error)
}
}
}
// Advanced tests section
if len(result.AdvancedTests) > 0 {
fmt.Println("\nAdvanced Tests:")
for _, test := range result.AdvancedTests {
status := "✓"
if !test.Success {
status = "✗"
}
name := test.Name
if test.Streaming {
name = fmt.Sprintf("Streaming %s", name)
}
fmt.Printf("[%s] %s (%.3fs)\n", status, name, float64(test.LatencyMs)/1000)
if !test.Success && test.Error != nil {
fmt.Printf(" Error: %v\n", test.Error)
}
}
} else if result.SkippedAdvanced {
fmt.Println("\nAdvanced Tests:")
fmt.Printf(" %s\n", result.SkippedReason)
}
// Summary
successRate := float64(result.TotalSuccess) / float64(result.TotalTests) * 100
fmt.Printf("\nSummary: %d/%d (%.2f%%) successful tests\n",
result.TotalSuccess, result.TotalTests, successRate)
fmt.Printf("Average latency: %.3fs\n", result.AverageLatency.Seconds())
}
// PrintSummaryReport prints the overall summary table of results
func PrintSummaryReport(results []AgentTestResult) {
fmt.Println("\nOverall Testing Summary:")
fmt.Println("=================================================")
// Create a tabwriter for aligned columns
w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
fmt.Fprintln(w, "Agent\tModel\tReasoning\tSuccess Rate\tAvg Latency\t")
fmt.Fprintln(w, "-----\t-----\t----------\t-----------\t-----------\t")
var totalSuccess, totalTests int
var totalLatency time.Duration
for _, result := range results {
success := result.TotalSuccess
total := result.TotalTests
successRate := float64(success) / float64(total) * 100
fmt.Fprintf(w, "%s\t%s\t%t\t%d/%d (%.2f%%)\t%.3fs\t\n",
result.AgentType,
result.ModelName,
result.Reasoning,
success,
total,
successRate,
result.AverageLatency.Seconds())
totalSuccess += success
totalTests += total
totalLatency += result.AverageLatency * time.Duration(total)
}
w.Flush()
if totalTests > 0 {
overallSuccessRate := float64(totalSuccess) / float64(totalTests) * 100
overallAvgLatency := totalLatency / time.Duration(totalTests)
fmt.Printf("\nTotal: %d/%d (%.2f%%) successful tests\n", totalSuccess, totalTests, overallSuccessRate)
fmt.Printf("Overall average latency: %.3fs\n", overallAvgLatency.Seconds())
}
}
// WriteReportToFile writes the test results to a report file in Markdown format
func WriteReportToFile(results []AgentTestResult, filePath string) error {
file, err := os.Create(filePath)
if err != nil {
return err
}
defer file.Close()
// Write header
file.WriteString("# LLM Agent Testing Report\n\n")
file.WriteString(fmt.Sprintf("Generated: %s\n\n", time.Now().UTC().Format(time.RFC1123)))
// Create a table for overall results
file.WriteString("## Overall Results\n\n")
file.WriteString("| Agent | Model | Reasoning | Success Rate | Average Latency |\n")
file.WriteString("|-------|-------|-----------|--------------|-----------------|\n")
var totalSuccess, totalTests int
var totalLatency time.Duration
for _, result := range results {
success := result.TotalSuccess
total := result.TotalTests
successRate := float64(success) / float64(total) * 100
file.WriteString(fmt.Sprintf("| %s | %s | %t | %d/%d (%.2f%%) | %.3fs |\n",
result.AgentType,
result.ModelName,
result.Reasoning,
success,
total,
successRate,
result.AverageLatency.Seconds()))
totalSuccess += success
totalTests += total
totalLatency += result.AverageLatency * time.Duration(total)
}
// Write summary
if totalTests > 0 {
overallSuccessRate := float64(totalSuccess) / float64(totalTests) * 100
overallAvgLatency := totalLatency / time.Duration(totalTests)
file.WriteString(fmt.Sprintf("\n**Total**: %d/%d (%.2f%%) successful tests\n",
totalSuccess, totalTests, overallSuccessRate))
file.WriteString(fmt.Sprintf("**Overall average latency**: %.3fs\n\n", overallAvgLatency.Seconds()))
}
// Write detailed results for each agent
file.WriteString("## Detailed Results\n\n")
for _, result := range results {
file.WriteString(fmt.Sprintf("### %s (%s)\n\n", result.AgentType, result.ModelName))
// Basic tests
if len(result.BasicTests) > 0 {
file.WriteString("#### Basic Tests\n\n")
file.WriteString("| Test | Result | Latency | Error |\n")
file.WriteString("|------|--------|---------|-------|\n")
for _, test := range result.BasicTests {
status := "✅ Pass"
errorMsg := ""
if !test.Success {
status = "❌ Fail"
if test.Error != nil {
errorMsg = TruncateString(EscapeMarkdown(test.Error.Error()), 150)
}
}
name := test.Name
if test.Streaming {
name = fmt.Sprintf("Streaming %s", name)
}
file.WriteString(fmt.Sprintf("| %s | %s | %.3fs | %s |\n",
name,
status,
float64(test.LatencyMs)/1000,
errorMsg))
}
file.WriteString("\n")
}
// Advanced tests
if len(result.AdvancedTests) > 0 {
file.WriteString("#### Advanced Tests\n\n")
file.WriteString("| Test | Result | Latency | Error |\n")
file.WriteString("|------|--------|---------|-------|\n")
for _, test := range result.AdvancedTests {
status := "✅ Pass"
errorMsg := ""
if !test.Success {
status = "❌ Fail"
if test.Error != nil {
errorMsg = TruncateString(EscapeMarkdown(test.Error.Error()), 150)
}
}
name := test.Name
if test.Streaming {
name = fmt.Sprintf("Streaming %s", name)
}
file.WriteString(fmt.Sprintf("| %s | %s | %.3fs | %s |\n",
name,
status,
float64(test.LatencyMs)/1000,
errorMsg))
}
file.WriteString("\n")
} else if result.SkippedAdvanced {
file.WriteString("#### Advanced Tests\n\n")
file.WriteString(fmt.Sprintf("*%s*\n\n", result.SkippedReason))
}
// Summary
successRate := float64(result.TotalSuccess) / float64(result.TotalTests) * 100
file.WriteString(fmt.Sprintf("**Summary**: %d/%d (%.2f%%) successful tests\n\n",
result.TotalSuccess, result.TotalTests, successRate))
file.WriteString(fmt.Sprintf("**Average latency**: %.3fs\n\n", result.AverageLatency.Seconds()))
file.WriteString("---\n\n")
}
return nil
}
backend/cmd/ctester/utils.go
+53
View File
@@ -0,0 +1,53 @@
package main
import (
"strings"
)
// Helper functions
// TruncateString truncates a string to a specified maximum length and adds ellipsis
func TruncateString(s string, maxLength int) string {
s = strings.Trim(s, "\n\r\t ")
s = strings.ReplaceAll(s, "\n", " ")
s = strings.ReplaceAll(s, "\r", " ")
s = strings.ReplaceAll(s, "\t", " ")
if len(s) <= maxLength {
return s
}
return s[:maxLength-3] + "..."
}
// EscapeMarkdown escapes special characters in markdown
func EscapeMarkdown(text string) string {
if text == "" {
return ""
}
replacements := []struct {
from string
to string
}{
{"|", "\\|"},
{"*", "\\*"},
{"_", "\\_"},
{"`", "\\`"},
{"#", "\\#"},
{"-", "\\-"},
{".", "\\."},
{"!", "\\!"},
{"(", "\\("},
{")", "\\)"},
{"[", "\\["},
{"]", "\\]"},
{"{", "\\{"},
{"}", "\\}"},
}
result := text
for _, r := range replacements {
result = strings.Replace(result, r.from, r.to, -1)
}
return result
}
backend/cmd/etester/flush.go
+42
View File
@@ -0,0 +1,42 @@
package main
import (
"fmt"
"os"
"pentagi/pkg/terminal"
)
// flush deletes all documents from the embedding store
func (t *Tester) flush() error {
terminal.Warning("This will delete ALL documents from the embedding store.")
response, err := terminal.GetYesNoInputContext(t.ctx, "Are you sure you want to continue?", os.Stdin)
if err != nil {
return fmt.Errorf("failed to get yes/no input: %w", err)
}
if !response {
terminal.Info("Operation cancelled.")
return nil
}
tx, err := t.conn.Begin(t.ctx)
if err != nil {
return fmt.Errorf("failed to start transaction: %w", err)
}
defer tx.Rollback(t.ctx)
result, err := tx.Exec(t.ctx, fmt.Sprintf("DELETE FROM %s", t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to delete documents: %w", err)
}
if err := tx.Commit(t.ctx); err != nil {
return fmt.Errorf("failed to commit transaction: %w", err)
}
rowsAffected := result.RowsAffected()
terminal.Success("\nSuccessfully deleted %d documents from the embedding store.", rowsAffected)
return nil
}
backend/cmd/etester/info.go
+203
View File
@@ -0,0 +1,203 @@
package main
import (
"database/sql"
"fmt"
"strings"
"pentagi/pkg/terminal"
)
// info displays statistics about the embedding database
func (t *Tester) info() error {
terminal.PrintHeader("Database Information:")
terminal.PrintThinSeparator()
// Get total document count
var docCount int
err := t.conn.QueryRow(t.ctx,
fmt.Sprintf("SELECT COUNT(*) FROM %s", t.embeddingTableName)).Scan(&docCount)
if err != nil {
return fmt.Errorf("failed to get document count: %w", err)
}
terminal.PrintKeyValueFormat("Total documents", "%d", docCount)
if docCount == 0 {
terminal.Info("No documents in the database.")
return nil
}
// Get average document size
var avgSize float64
err = t.conn.QueryRow(t.ctx,
fmt.Sprintf("SELECT AVG(LENGTH(document)) FROM %s", t.embeddingTableName)).Scan(&avgSize)
if err != nil {
return fmt.Errorf("failed to get average document size: %w", err)
}
terminal.PrintKeyValueFormat("Average document size", "%.2f bytes", avgSize)
// Get total document size
var totalSize int64
err = t.conn.QueryRow(t.ctx,
fmt.Sprintf("SELECT SUM(LENGTH(document)) FROM %s", t.embeddingTableName)).Scan(&totalSize)
if err != nil {
return fmt.Errorf("failed to get total document size: %w", err)
}
terminal.PrintKeyValue("Total document size", formatSize(totalSize))
// Get document type distribution
terminal.PrintHeader("\nDocument Type Distribution:")
rows, err := t.conn.Query(t.ctx,
fmt.Sprintf("SELECT cmetadata->>'doc_type' as type, COUNT(*) FROM %s GROUP BY type ORDER BY COUNT(*) DESC",
t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to get document type distribution: %w", err)
}
defer rows.Close()
printTableHeader("Type", "Count")
for rows.Next() {
var docType sql.NullString
var count int
if err := rows.Scan(&docType, &count); err != nil {
return fmt.Errorf("failed to scan document type row: %w", err)
}
typeStr := "unknown"
if docType.Valid {
typeStr = docType.String
}
printTableRow(typeStr, count)
}
// Get flow_id distribution
terminal.PrintHeader("\nFlow ID Distribution:")
rows, err = t.conn.Query(t.ctx,
fmt.Sprintf("SELECT cmetadata->>'flow_id' as flow_id, COUNT(*) FROM %s GROUP BY flow_id ORDER BY COUNT(*) DESC",
t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to get flow ID distribution: %w", err)
}
defer rows.Close()
printTableHeader("Flow ID", "Count")
for rows.Next() {
var flowID sql.NullString
var count int
if err := rows.Scan(&flowID, &count); err != nil {
return fmt.Errorf("failed to scan flow ID row: %w", err)
}
flowStr := "unknown"
if flowID.Valid {
flowStr = flowID.String
}
printTableRow(flowStr, count)
}
// Get guide_type distribution for doc_type = 'guide'
terminal.PrintHeader("\nGuide Type Distribution (for doc_type = 'guide'):")
rows, err = t.conn.Query(t.ctx,
fmt.Sprintf("SELECT cmetadata->>'guide_type' as guide_type, COUNT(*) FROM %s "+
"WHERE cmetadata->>'doc_type' = 'guide' GROUP BY guide_type ORDER BY COUNT(*) DESC",
t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to get guide type distribution: %w", err)
}
defer rows.Close()
printTableHeader("Guide Type", "Count")
hasRows := false
for rows.Next() {
hasRows = true
var guideType sql.NullString
var count int
if err := rows.Scan(&guideType, &count); err != nil {
return fmt.Errorf("failed to scan guide type row: %w", err)
}
typeStr := "unknown"
if guideType.Valid {
typeStr = guideType.String
}
printTableRow(typeStr, count)
}
if !hasRows {
terminal.Info("No guide documents found.")
}
// Get code_lang distribution for doc_type = 'code'
terminal.PrintHeader("\nCode Language Distribution (for doc_type = 'code'):")
rows, err = t.conn.Query(t.ctx,
fmt.Sprintf("SELECT cmetadata->>'code_lang' as code_lang, COUNT(*) FROM %s "+
"WHERE cmetadata->>'doc_type' = 'code' GROUP BY code_lang ORDER BY COUNT(*) DESC",
t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to get code language distribution: %w", err)
}
defer rows.Close()
printTableHeader("Code Language", "Count")
hasRows = false
for rows.Next() {
hasRows = true
var codeLang sql.NullString
var count int
if err := rows.Scan(&codeLang, &count); err != nil {
return fmt.Errorf("failed to scan code language row: %w", err)
}
langStr := "unknown"
if codeLang.Valid {
langStr = codeLang.String
}
printTableRow(langStr, count)
}
if !hasRows {
terminal.Info("No code documents found.")
}
// Get answer_type distribution for doc_type = 'answer'
terminal.PrintHeader("\nAnswer Type Distribution (for doc_type = 'answer'):")
rows, err = t.conn.Query(t.ctx,
fmt.Sprintf("SELECT cmetadata->>'answer_type' as answer_type, COUNT(*) FROM %s "+
"WHERE cmetadata->>'doc_type' = 'answer' GROUP BY answer_type ORDER BY COUNT(*) DESC",
t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to get answer type distribution: %w", err)
}
defer rows.Close()
printTableHeader("Answer Type", "Count")
hasRows = false
for rows.Next() {
hasRows = true
var answerType sql.NullString
var count int
if err := rows.Scan(&answerType, &count); err != nil {
return fmt.Errorf("failed to scan answer type row: %w", err)
}
typeStr := "unknown"
if answerType.Valid {
typeStr = answerType.String
}
printTableRow(typeStr, count)
}
if !hasRows {
terminal.Info("No answer documents found.")
}
return nil
}
// printTableHeader prints a formatted table header row
func printTableHeader(column1, column2 string) {
fmt.Printf("%-20s | %s\n", column1, column2)
fmt.Printf("%-20s-+-%s\n", strings.Repeat("-", 20), strings.Repeat("-", 10))
}
// printTableRow prints a table row with data
func printTableRow(value string, count int) {
fmt.Printf("%-20s | %d\n", value, count)
}
backend/cmd/etester/main.go
+135
View File
@@ -0,0 +1,135 @@
package main
import (
"context"
"flag"
"log"
"os"
"os/signal"
"syscall"
"time"
"pentagi/pkg/config"
"pentagi/pkg/providers/embeddings"
"pentagi/pkg/terminal"
"pentagi/pkg/version"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/joho/godotenv"
"github.com/sirupsen/logrus"
)
const (
defaultEmbeddingTableName = "langchain_pg_embedding"
defaultCollectionTableName = "langchain_pg_collection"
)
func main() {
// Define flags (but don't include command as a flag)
verbose := flag.Bool("verbose", false, "Enable verbose output")
envFile := flag.String("env", ".env", "Path to environment file")
help := flag.Bool("help", false, "Show help information")
flag.Parse()
logrus.Infof("Starting PentAGI Embedding Tester %s", version.GetBinaryVersion())
// Extract command from first non-flag argument
args := flag.Args()
var command string
if len(args) > 0 {
command = args[0]
args = args[1:] // Remove command from args
} else {
command = "test" // Default command
}
if *help {
showHelp()
return
}
// Load environment from .env file
err := godotenv.Load(*envFile)
if err != nil {
log.Println("Warning: Error loading .env file:", err)
}
cfg, err := config.NewConfig()
if err != nil {
log.Fatalf("Error loading config: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Initialize database connection pool
poolConfig, err := pgxpool.ParseConfig(cfg.DatabaseURL)
if err != nil {
log.Fatalf("Unable to parse database URL: %v", err)
}
poolConfig.MaxConns = 10
poolConfig.MinConns = 2
poolConfig.MaxConnLifetime = time.Hour
poolConfig.MaxConnIdleTime = 30 * time.Minute
connPool, err := pgxpool.NewWithConfig(ctx, poolConfig)
if err != nil {
log.Fatalf("Unable to create connection pool: %v", err)
}
defer connPool.Close()
embedder, err := embeddings.New(cfg)
if err != nil {
log.Fatalf("Unable to create embedder: %v", err)
}
// Initialize tester with the parsed command
tester := NewTester(
connPool,
embedder,
*verbose,
command,
ctx,
cfg,
)
// Handle graceful shutdown
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
go func() {
<-sigChan
terminal.Info("Shutting down gracefully...")
cancel()
}()
// Execute the command with remaining arguments
if err := tester.executeCommand(args); err != nil {
terminal.Error("Error executing command: %v", err)
os.Exit(1)
}
}
func showHelp() {
terminal.PrintHeader("Embedding Tester (etester) - A tool for testing and managing embeddings")
terminal.Info("\nUsage:")
terminal.Info(" ./etester [flags] [command] [args]")
terminal.Info("\nFlags:")
terminal.Info(" -env string Path to environment file (default \".env\")")
terminal.Info(" -verbose Enable verbose output")
terminal.Info(" -help Show this help message")
terminal.Info("\nCommands:")
terminal.PrintKeyValue(" test ", "Test embedding provider and pgvector connection")
terminal.PrintKeyValue(" info ", "Display statistics about the embedding database")
terminal.PrintKeyValue(" flush ", "Delete all documents from the embedding database")
terminal.PrintKeyValue(" reindex ", "Recalculate embeddings for all documents")
terminal.PrintKeyValue(" search ", "Search for documents in the embedding database")
terminal.Info("\nExamples:")
terminal.Info(" ./etester test -verbose Test with verbose output")
terminal.Info(" ./etester info Show database statistics")
terminal.Info(" ./etester flush Delete all documents")
terminal.Info(" ./etester reindex Reindex all documents")
terminal.Info(" ./etester search -query \"How to install PostgreSQL\" Search for documents")
terminal.Info("")
}
backend/cmd/etester/reindex.go
+122
View File
@@ -0,0 +1,122 @@
package main
import (
"fmt"
"os"
"pentagi/pkg/terminal"
"github.com/jackc/pgx/v5"
"github.com/pgvector/pgvector-go"
)
// Document represents a document in the embedding store
type Document struct {
UUID string
Content string
}
// reindex recalculates embeddings for all documents in the store
func (t *Tester) reindex() error {
terminal.Warning("This will reindex ALL documents in the embedding store.")
terminal.Warning("This operation may take a long time depending on the number of documents.")
response, err := terminal.GetYesNoInputContext(t.ctx, "Are you sure you want to continue?", os.Stdin)
if err != nil {
return fmt.Errorf("failed to get yes/no input: %w", err)
}
if !response {
terminal.Info("Operation cancelled.")
return nil
}
// Get total document count
var totalDocs int
err = t.conn.QueryRow(t.ctx, fmt.Sprintf("SELECT COUNT(*) FROM %s", t.embeddingTableName)).Scan(&totalDocs)
if err != nil {
return fmt.Errorf("failed to get document count: %w", err)
}
if totalDocs == 0 {
terminal.Info("No documents found in the embedding store.")
return nil
}
terminal.Info(fmt.Sprintf("Found %d documents to reindex.", totalDocs))
// Calculate batch size for processing
batchSize := t.cfg.EmbeddingBatchSize
if batchSize <= 0 {
batchSize = 10 // Default batch size
}
rows, err := t.conn.Query(t.ctx, fmt.Sprintf("SELECT uuid, document FROM %s", t.embeddingTableName))
if err != nil {
return fmt.Errorf("failed to query documents: %w", err)
}
defer rows.Close()
// Collect documents
documents := []Document{}
for rows.Next() {
var doc Document
if err := rows.Scan(&doc.UUID, &doc.Content); err != nil {
return fmt.Errorf("failed to scan document row: %w", err)
}
documents = append(documents, doc)
}
if err := rows.Err(); err != nil {
return fmt.Errorf("error iterating document rows: %w", err)
}
totalBatches := (len(documents) + batchSize - 1) / batchSize
processedDocs := 0
// Process documents in batches to avoid memory issues
for i := 0; i < totalBatches; i++ {
start := i * batchSize
end := min((i+1)*batchSize, len(documents))
batchDocs := documents[start:end]
// Extract content for embedding
texts := make([]string, len(batchDocs))
for j, doc := range batchDocs {
texts[j] = doc.Content
}
// Generate embeddings
terminal.Info(fmt.Sprintf("Processing batch %d/%d (%d documents)...",
i+1, totalBatches, len(batchDocs)))
vectors, err := t.embedder.EmbedDocuments(t.ctx, texts)
if err != nil {
return fmt.Errorf("failed to generate embeddings for batch %d: %w", i+1, err)
}
if len(vectors) != len(batchDocs) {
return fmt.Errorf("embedder returned wrong number of vectors: got %d, expected %d",
len(vectors), len(batchDocs))
}
// Update documents in database
batch := &pgx.Batch{}
for j, doc := range batchDocs {
batch.Queue(
fmt.Sprintf("UPDATE %s SET embedding = $1 WHERE uuid = $2", t.embeddingTableName),
pgvector.NewVector(vectors[j]), doc.UUID)
}
results := t.conn.SendBatch(t.ctx, batch)
if err := results.Close(); err != nil {
return fmt.Errorf("failed to update embeddings for batch %d: %w", i+1, err)
}
processedDocs += len(batchDocs)
progressPercent := float64(processedDocs) / float64(totalDocs) * 100
terminal.Info("Progress: %.2f%% (%d/%d documents processed)", progressPercent, processedDocs, totalDocs)
}
terminal.Success("\nReindexing completed successfully! %d documents were updated.", processedDocs)
return nil
}
backend/cmd/etester/search.go
+289
View File
@@ -0,0 +1,289 @@
package main
import (
"fmt"
"sort"
"strconv"
"strings"
"pentagi/pkg/terminal"
"github.com/vxcontrol/langchaingo/vectorstores"
"github.com/vxcontrol/langchaingo/vectorstores/pgvector"
)
// SearchOptions represents the options for vector search
type SearchOptions struct {
Query string
DocType string
FlowID int64
AnswerType string
GuideType string
Limit int
Threshold float32
}
// Validates and fills in default values for search options
func validateSearchOptions(opts *SearchOptions) error {
// Query is required
if opts.Query == "" {
return fmt.Errorf("query parameter is required")
}
// Validate doc_type if provided
if opts.DocType != "" {
validDocTypes := map[string]bool{
"answer": true,
"memory": true,
"guide": true,
"code": true,
}
if !validDocTypes[opts.DocType] {
return fmt.Errorf("invalid doc_type: %s. Valid values are: answer, memory, guide, code", opts.DocType)
}
}
// Validate flow_id if provided
if opts.FlowID < 0 {
return fmt.Errorf("flow_id must be a positive number")
}
// Validate answer_type if provided
if opts.AnswerType != "" {
validAnswerTypes := map[string]bool{
"guide": true,
"vulnerability": true,
"code": true,
"tool": true,
"other": true,
}
if !validAnswerTypes[opts.AnswerType] {
return fmt.Errorf("invalid answer_type: %s. Valid values are: guide, vulnerability, code, tool, other", opts.AnswerType)
}
}
// Validate guide_type if provided
if opts.GuideType != "" {
validGuideTypes := map[string]bool{
"install": true,
"configure": true,
"use": true,
"pentest": true,
"development": true,
"other": true,
}
if !validGuideTypes[opts.GuideType] {
return fmt.Errorf("invalid guide_type: %s. Valid values are: install, configure, use, pentest, development, other", opts.GuideType)
}
}
// Validate limit
if opts.Limit <= 0 {
opts.Limit = 3 // Default limit
}
// Validate threshold
if opts.Threshold <= 0 || opts.Threshold > 1 {
opts.Threshold = 0.7 // Default threshold
}
return nil
}
// ParseSearchArgs parses command line arguments specific for search
func parseSearchArgs(args []string) (*SearchOptions, error) {
if len(args) == 0 {
return nil, fmt.Errorf("no arguments provided")
}
opts := &SearchOptions{}
for i := 0; i < len(args); i++ {
arg := args[i]
if !strings.HasPrefix(arg, "-") {
continue
}
paramName := strings.TrimPrefix(arg, "-")
if i+1 >= len(args) || strings.HasPrefix(args[i+1], "-") {
return nil, fmt.Errorf("missing value for parameter: %s", paramName)
}
paramValue := args[i+1]
i++
switch paramName {
case "query":
opts.Query = paramValue
case "doc_type":
opts.DocType = paramValue
case "flow_id":
flowID, err := strconv.ParseInt(paramValue, 10, 64)
if err != nil {
return nil, fmt.Errorf("invalid flow_id value: %v", err)
}
opts.FlowID = flowID
case "answer_type":
opts.AnswerType = paramValue
case "guide_type":
opts.GuideType = paramValue
case "limit":
limit, err := strconv.Atoi(paramValue)
if err != nil {
return nil, fmt.Errorf("invalid limit value: %v", err)
}
opts.Limit = limit
case "threshold":
threshold, err := strconv.ParseFloat(paramValue, 32)
if err != nil {
return nil, fmt.Errorf("invalid threshold value: %v", err)
}
opts.Threshold = float32(threshold)
default:
return nil, fmt.Errorf("unknown parameter: %s", paramName)
}
}
if err := validateSearchOptions(opts); err != nil {
return nil, err
}
return opts, nil
}
// search performs vector search in the embedding database
func (t *Tester) search(args []string) error {
// Display usage if no arguments provided
if len(args) == 0 {
printSearchUsage()
return nil
}
// Parse search options
opts, err := parseSearchArgs(args)
if err != nil {
terminal.Error("Error parsing search arguments: %v", err)
printSearchUsage()
return nil
}
// Create pgvector store if needed for search
store, err := t.createVectorStore()
if err != nil {
return fmt.Errorf("failed to create vector store: %w", err)
}
// Prepare filters
filters := make(map[string]any)
if opts.DocType != "" {
filters["doc_type"] = opts.DocType
}
if opts.FlowID > 0 {
filters["flow_id"] = strconv.FormatInt(opts.FlowID, 10)
}
if opts.AnswerType != "" {
filters["answer_type"] = opts.AnswerType
}
if opts.GuideType != "" {
filters["guide_type"] = opts.GuideType
}
// Prepare search options
searchOpts := []vectorstores.Option{
vectorstores.WithScoreThreshold(opts.Threshold),
}
if len(filters) > 0 {
searchOpts = append(searchOpts, vectorstores.WithFilters(filters))
}
// Perform the search
terminal.Info("Searching for: %s", opts.Query)
terminal.Info("Threshold: %.2f, Limit: %d", opts.Threshold, opts.Limit)
if len(filters) > 0 {
terminal.Info("Filters: %v", filters)
}
docs, err := store.SimilaritySearch(
t.ctx,
opts.Query,
opts.Limit,
searchOpts...,
)
if err != nil {
return fmt.Errorf("search failed: %w", err)
}
// Display results
if len(docs) == 0 {
terminal.Info("No matching documents found.")
return nil
}
terminal.Success("Found %d matching documents:", len(docs))
terminal.PrintThinSeparator()
for i, doc := range docs {
terminal.PrintHeader(fmt.Sprintf("Result #%d (similarity score: %.4f)", i+1, doc.Score))
// Print metadata
terminal.Info("Metadata:")
keys := []string{}
for k := range doc.Metadata {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
terminal.PrintKeyValueFormat(fmt.Sprintf("%-12s ", k), "%v", doc.Metadata[k])
}
// Print content with markdown rendering
terminal.PrintThinSeparator()
terminal.PrintResult(doc.PageContent)
terminal.PrintThickSeparator()
}
return nil
}
// createVectorStore creates a pgvector store instance using the current connection and embedder
func (t *Tester) createVectorStore() (*pgvector.Store, error) {
// Create pgvector store
store, err := pgvector.New(
t.ctx,
pgvector.WithConn(t.conn),
pgvector.WithEmbedder(t.embedder),
pgvector.WithCollectionName("langchain"),
pgvector.WithEmbeddingTableName(t.embeddingTableName),
pgvector.WithCollectionTableName(t.collectionTableName),
)
if err != nil {
return nil, err
}
return &store, nil
}
// printSearchUsage prints the usage information for the search command
func printSearchUsage() {
terminal.PrintHeader("Search Command Usage:")
terminal.Info("Performs vector search in the embedding database")
terminal.Info("\nSyntax:")
terminal.Info(" ./etester search [OPTIONS]")
terminal.Info("\nOptions:")
terminal.PrintKeyValue(" -query STRING", "Search query text (required)")
terminal.PrintKeyValue(" -doc_type STRING", "Filter by document type (answer, memory, guide, code)")
terminal.PrintKeyValue(" -flow_id NUMBER", "Filter by flow ID (positive number)")
terminal.PrintKeyValue(" -answer_type STRING", "Filter by answer type (guide, vulnerability, code, tool, other)")
terminal.PrintKeyValue(" -guide_type STRING", "Filter by guide type (install, configure, use, pentest, development, other)")
terminal.PrintKeyValue(" -limit NUMBER", "Maximum number of results (default: 3)")
terminal.PrintKeyValue(" -threshold NUMBER", "Similarity threshold (0.0-1.0, default: 0.7)")
terminal.Info("\nExamples:")
terminal.Info(" ./etester search -query \"How to install PostgreSQL\" -limit 5")
terminal.Info(" ./etester search -query \"Security vulnerability\" -doc_type guide -threshold 0.8")
terminal.Info(" ./etester search -query \"Code examples\" -doc_type code -flow_id 42")
}
backend/cmd/etester/test.go
+106
View File
@@ -0,0 +1,106 @@
package main
import (
"database/sql"
"fmt"
"strings"
"pentagi/pkg/terminal"
)
const (
testText = "This is a test text for embedding"
testTexts = "This is a test text for embedding\nThis is another test text for embedding"
)
// test checks connectivity to the database and tests the embedder.
func (t *Tester) test() error {
terminal.Info("Testing connection to PostgreSQL database... ")
err := t.conn.Ping(t.ctx)
if err != nil {
terminal.Error("FAILED")
return fmt.Errorf("database connection test failed: %w", err)
}
terminal.Success("OK")
terminal.Info("Testing pgvector extension... ")
var result string
err = t.conn.QueryRow(t.ctx, "SELECT extname FROM pg_extension WHERE extname = 'vector'").Scan(&result)
if err != nil {
if err == sql.ErrNoRows {
terminal.Error("FAILED")
return fmt.Errorf("pgvector extension is not installed")
}
terminal.Error("FAILED")
return fmt.Errorf("failed to check pgvector extension: %w", err)
}
terminal.Success("OK")
terminal.Info("Testing embedding table existence... ")
var tableExists bool
err = t.conn.QueryRow(t.ctx,
"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = $1)",
t.embeddingTableName).Scan(&tableExists)
if err != nil {
terminal.Error("FAILED")
return fmt.Errorf("failed to check embedding table: %w", err)
}
if !tableExists {
terminal.Error("FAILED")
return fmt.Errorf("embedding table '%s' does not exist", t.embeddingTableName)
}
terminal.Success("OK")
terminal.Info("Testing embedder with single query... ")
if !t.embedder.IsAvailable() {
terminal.Error("FAILED")
return fmt.Errorf("embedder is not available")
}
embedVector, err := t.embedder.EmbedQuery(t.ctx, testText)
if err != nil {
terminal.Error("FAILED")
return fmt.Errorf("embedder test failed: %w", err)
}
if len(embedVector) == 0 {
terminal.Error("FAILED")
return fmt.Errorf("embedder returned empty vector")
}
terminal.Success(fmt.Sprintf("OK (%d dimensions)", len(embedVector)))
terminal.Info("Testing embedder with multiple documents... ")
texts := strings.Split(testTexts, "\n")
embedVectors, err := t.embedder.EmbedDocuments(t.ctx, texts)
if err != nil {
terminal.Error("FAILED")
return fmt.Errorf("embedder multi-text test failed: %w", err)
}
if len(embedVectors) != len(texts) {
terminal.Error("FAILED")
return fmt.Errorf("embedder returned wrong number of vectors: got %d, expected %d",
len(embedVectors), len(texts))
}
if len(embedVectors[0]) == 0 || len(embedVectors[1]) == 0 {
terminal.Error("FAILED")
return fmt.Errorf("embedder returned empty vectors")
}
terminal.Success(fmt.Sprintf("OK (%d documents, %d dimensions each)",
len(embedVectors), len(embedVectors[0])))
if t.verbose {
terminal.PrintHeader("\nVerbose output:")
terminal.PrintKeyValue("Embedding provider", t.cfg.EmbeddingProvider)
terminal.PrintKeyValueFormat("Vector dimensions", "%d", len(embedVector))
// Display a sample of vector values for inspection
vectorSample := embedVector
if len(vectorSample) > 5 {
vectorSample = vectorSample[:5]
}
terminal.PrintKeyValue("First values of test vector",
fmt.Sprintf("%v", vectorSample))
}
terminal.Success("\nAll tests passed successfully!")
return nil
}
backend/cmd/etester/tester.go
+76
View File
@@ -0,0 +1,76 @@
package main
import (
"context"
"fmt"
"pentagi/pkg/config"
"pentagi/pkg/providers/embeddings"
"github.com/jackc/pgx/v5/pgxpool"
)
// Tester represents the main application structure for the etester tool
type Tester struct {
conn *pgxpool.Pool
embedder embeddings.Embedder
embeddingTableName string
collectionTableName string
verbose bool
command string
ctx context.Context
cfg *config.Config
}
// NewTester creates a new instance of the Tester with the provided configuration
func NewTester(
conn *pgxpool.Pool,
embedder embeddings.Embedder,
verbose bool,
command string,
ctx context.Context,
cfg *config.Config,
) *Tester {
return &Tester{
conn: conn,
embedder: embedder,
embeddingTableName: defaultEmbeddingTableName,
collectionTableName: defaultCollectionTableName,
verbose: verbose,
command: command,
ctx: ctx,
cfg: cfg,
}
}
// executeCommand executes the appropriate command based on the command string
func (t *Tester) executeCommand(args []string) error {
switch t.command {
case "test":
return t.test()
case "info":
return t.info()
case "flush":
return t.flush()
case "reindex":
return t.reindex()
case "search":
return t.search(args)
default:
return fmt.Errorf("unknown command: %s", t.command)
}
}
// formatSize formats a file size in bytes to a human-readable string
func formatSize(bytes int64) string {
const unit = 1024
if bytes < unit {
return fmt.Sprintf("%d B", bytes)
}
div, exp := int64(unit), 0
for n := bytes / unit; n >= unit; n /= unit {
div *= unit
exp++
}
return fmt.Sprintf("%.2f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}
backend/cmd/ftester/main.go
+162
View File
@@ -0,0 +1,162 @@
package main
import (
"context"
"database/sql"
"errors"
"flag"
"fmt"
"log"
"os"
"os/signal"
"syscall"
"time"
"pentagi/cmd/ftester/worker"
"pentagi/pkg/config"
"pentagi/pkg/database"
"pentagi/pkg/docker"
obs "pentagi/pkg/observability"
"pentagi/pkg/providers"
"pentagi/pkg/providers/provider"
"pentagi/pkg/terminal"
"pentagi/pkg/version"
"github.com/joho/godotenv"
_ "github.com/lib/pq"
"github.com/sirupsen/logrus"
)
func main() {
envFile := flag.String("env", ".env", "Path to environment file")
providerName := flag.String("provider", "custom", "Provider name (openai, anthropic, gemini, bedrock, ollama, deepseek, glm, kimi, qwen, custom)")
flowID := flag.Int64("flow", 0, "Flow ID for testing functions that require it (0 means using mocks)")
userID := flag.Int64("user", 0, "User ID for testing functions that require it (1 is default admin user)")
taskID := flag.Int64("task", 0, "Task ID for testing functions with default unset")
subtaskID := flag.Int64("subtask", 0, "Subtask ID for testing functions with default unset")
flag.Parse()
if *taskID == 0 {
taskID = nil
}
if *subtaskID == 0 {
subtaskID = nil
}
logrus.Infof("Starting PentAGI Function Tester %s", version.GetBinaryVersion())
err := godotenv.Load(*envFile)
if err != nil {
log.Println("Warning: Error loading .env file:", err)
}
cfg, err := config.NewConfig()
if err != nil {
log.Fatalf("Error loading config: %v", err)
}
// Setup signal handling for graceful shutdown
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
lfclient, err := obs.NewLangfuseClient(ctx, cfg)
if err != nil && !errors.Is(err, obs.ErrNotConfigured) {
log.Fatalf("Unable to create langfuse client: %v\n", err)
}
defer func() {
if lfclient != nil {
lfclient.ForceFlush(context.Background())
}
}()
otelclient, err := obs.NewTelemetryClient(ctx, cfg)
if err != nil && !errors.Is(err, obs.ErrNotConfigured) {
log.Fatalf("Unable to create telemetry client: %v\n", err)
}
defer func() {
if otelclient != nil {
otelclient.ForceFlush(context.Background())
}
}()
obs.InitObserver(ctx, lfclient, otelclient, []logrus.Level{
logrus.DebugLevel,
logrus.InfoLevel,
logrus.WarnLevel,
logrus.ErrorLevel,
})
// Initialize database connection
db, err := sql.Open("postgres", cfg.DatabaseURL)
if err != nil {
log.Fatalf("Unable to open database: %v", err)
}
db.SetMaxOpenConns(10)
db.SetMaxIdleConns(2)
db.SetConnMaxLifetime(time.Hour)
queries := database.New(db)
terminal.PrintHeader("Function Tester (ftester)")
terminal.PrintInfo("Starting ftester with the following parameters:")
terminal.PrintKeyValue("Environment file", *envFile)
terminal.PrintKeyValue("Provider", *providerName)
if *flowID != 0 {
terminal.PrintKeyValue("Flow ID", fmt.Sprintf("%d", *flowID))
} else {
terminal.PrintInfo("Using mock mode (flowID=0)")
}
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
terminal.PrintThinSeparator()
// Initialize docker client
dockerClient, err := docker.NewDockerClient(context.Background(), queries, cfg)
if err != nil {
log.Fatalf("Failed to initialize Docker client: %v", err)
}
// Initialize provider controller
providerController, err := providers.NewProviderController(cfg, queries, dockerClient)
if err != nil {
log.Fatalf("Failed to initialize provider controller: %v", err)
}
// Initialize tester with appropriate proxy interfaces
tester, err := worker.NewTester(
queries,
cfg,
ctx,
dockerClient,
providerController,
*flowID,
*userID,
taskID,
subtaskID,
provider.ProviderName(*providerName),
)
if err != nil {
log.Fatalf("Failed to initialize tester worker: %v", err)
}
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
go func() {
<-sigChan
fmt.Println("\nShutting down gracefully...")
cancel()
}()
// Execute the tester with the parsed arguments
if err := tester.Execute(flag.Args()); err != nil {
terminal.PrintError("Error executing function: %v", err)
os.Exit(1)
}
}
backend/cmd/ftester/mocks/logs.go
+352
View File
@@ -0,0 +1,352 @@
package mocks
import (
"context"
"encoding/json"
"fmt"
"pentagi/pkg/database"
"pentagi/pkg/graph/model"
"pentagi/pkg/terminal"
"pentagi/pkg/tools"
)
type ProxyProviders interface {
GetScreenshotProvider() tools.ScreenshotProvider
GetAgentLogProvider() tools.AgentLogProvider
GetMsgLogProvider() tools.MsgLogProvider
GetSearchLogProvider() tools.SearchLogProvider
GetTermLogProvider() tools.TermLogProvider
GetVectorStoreLogProvider() tools.VectorStoreLogProvider
GetToolCallLogProvider() tools.ToolCallLogProvider
GetKnowledgeProvider() tools.KnowledgeProvider
}
// proxyProviders contains all the proxy implementations for various providers
type proxyProviders struct {
screenshot *proxyScreenshotProvider
agentLog *proxyAgentLogProvider
msgLog *proxyMsgLogProvider
searchLog *proxySearchLogProvider
termLog *proxyTermLogProvider
vectorStoreLog *proxyVectorStoreLogProvider
toolCallLog *proxyToolCallLogProvider
knowledgeProvider *proxyKnowledgeProvider
}
// NewProxyProviders creates a new set of proxy providers
func NewProxyProviders() ProxyProviders {
return &proxyProviders{
screenshot: &proxyScreenshotProvider{},
agentLog: &proxyAgentLogProvider{},
msgLog: &proxyMsgLogProvider{},
searchLog: &proxySearchLogProvider{},
termLog: &proxyTermLogProvider{},
vectorStoreLog: &proxyVectorStoreLogProvider{},
toolCallLog: &proxyToolCallLogProvider{},
knowledgeProvider: &proxyKnowledgeProvider{},
}
}
func (p *proxyProviders) GetScreenshotProvider() tools.ScreenshotProvider {
return p.screenshot
}
func (p *proxyProviders) GetAgentLogProvider() tools.AgentLogProvider {
return p.agentLog
}
func (p *proxyProviders) GetMsgLogProvider() tools.MsgLogProvider {
return p.msgLog
}
func (p *proxyProviders) GetSearchLogProvider() tools.SearchLogProvider {
return p.searchLog
}
func (p *proxyProviders) GetTermLogProvider() tools.TermLogProvider {
return p.termLog
}
func (p *proxyProviders) GetVectorStoreLogProvider() tools.VectorStoreLogProvider {
return p.vectorStoreLog
}
func (p *proxyProviders) GetToolCallLogProvider() tools.ToolCallLogProvider {
return p.toolCallLog
}
func (p *proxyProviders) GetKnowledgeProvider() tools.KnowledgeProvider {
return p.knowledgeProvider
}
// proxyScreenshotProvider is a proxy implementation of ScreenshotProvider
type proxyScreenshotProvider struct{}
// PutScreenshot implements the ScreenshotProvider interface
func (p *proxyScreenshotProvider) PutScreenshot(ctx context.Context, name, url string, taskID, subtaskID *int64) (int64, error) {
terminal.PrintInfo("Screenshot saved:")
terminal.PrintKeyValue("Name", name)
terminal.PrintKeyValue("URL", url)
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
return 0, nil
}
// proxyAgentLogProvider is a proxy implementation of AgentLogProvider
type proxyAgentLogProvider struct{}
// PutLog implements the AgentLogProvider interface
func (p *proxyAgentLogProvider) PutLog(
ctx context.Context,
initiator database.MsgchainType,
executor database.MsgchainType,
task string,
result string,
taskID *int64,
subtaskID *int64,
) (int64, error) {
terminal.PrintInfo("Agent log saved:")
terminal.PrintKeyValue("Initiator", string(initiator))
terminal.PrintKeyValue("Executor", string(executor))
terminal.PrintKeyValue("Task", task)
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
if len(result) > 0 {
terminal.PrintResultWithKey("Result", result)
}
return 0, nil
}
// proxyMsgLogProvider is a proxy implementation of MsgLogProvider
type proxyMsgLogProvider struct{}
// PutMsg implements the MsgLogProvider interface
func (p *proxyMsgLogProvider) PutMsg(
ctx context.Context,
msgType database.MsglogType,
taskID, subtaskID *int64,
streamID int64, // unsupported for now
thinking, msg string,
) (int64, error) {
terminal.PrintInfo("Message logged:")
terminal.PrintKeyValue("Type", string(msgType))
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
if len(msg) > 0 {
terminal.PrintResultWithKey("Message", msg)
}
return 0, nil
}
// UpdateMsgResult implements the MsgLogProvider interface
func (p *proxyMsgLogProvider) UpdateMsgResult(
ctx context.Context,
msgID int64,
streamID int64, // unsupported for now
result string,
resultFormat database.MsglogResultFormat,
) error {
terminal.PrintInfo("Message result updated:")
terminal.PrintKeyValueFormat("Message ID", "%d", msgID)
terminal.PrintKeyValue("Format", string(resultFormat))
if len(result) > 0 {
terminal.PrintResultWithKey("Result", result)
}
return nil
}
// proxySearchLogProvider is a proxy implementation of SearchLogProvider
type proxySearchLogProvider struct{}
// PutLog implements the SearchLogProvider interface
func (p *proxySearchLogProvider) PutLog(
ctx context.Context,
initiator database.MsgchainType,
executor database.MsgchainType,
engine database.SearchengineType,
query string,
result string,
taskID *int64,
subtaskID *int64,
) (int64, error) {
terminal.PrintInfo("Search log saved:")
terminal.PrintKeyValue("Initiator", string(initiator))
terminal.PrintKeyValue("Executor", string(executor))
terminal.PrintKeyValue("Engine", string(engine))
terminal.PrintKeyValue("Query", query)
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
if len(result) > 0 {
terminal.PrintResultWithKey("Search Result", result)
}
return 0, nil
}
// proxyTermLogProvider is a proxy implementation of TermLogProvider
type proxyTermLogProvider struct{}
// PutMsg implements the TermLogProvider interface
func (p *proxyTermLogProvider) PutMsg(
ctx context.Context,
msgType database.TermlogType,
msg string,
containerID int64,
taskID, subtaskID *int64,
) (int64, error) {
terminal.PrintInfo("Terminal log saved:")
terminal.PrintKeyValue("Type", string(msgType))
terminal.PrintKeyValueFormat("Container ID", "%d", containerID)
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
if len(msg) > 0 {
terminal.PrintResultWithKey("Terminal Output", msg)
}
return 0, nil
}
// proxyVectorStoreLogProvider is a proxy implementation of VectorStoreLogProvider
type proxyVectorStoreLogProvider struct{}
// PutLog implements the VectorStoreLogProvider interface
func (p *proxyVectorStoreLogProvider) PutLog(
ctx context.Context,
initiator database.MsgchainType,
executor database.MsgchainType,
filter json.RawMessage,
query string,
action database.VecstoreActionType,
result string,
taskID *int64,
subtaskID *int64,
) (int64, error) {
terminal.PrintInfo("Vector store log saved:")
terminal.PrintKeyValue("Initiator", string(initiator))
terminal.PrintKeyValue("Executor", string(executor))
terminal.PrintKeyValue("Action", string(action))
terminal.PrintKeyValue("Query", query)
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
if len(result) > 0 {
terminal.PrintResultWithKey("Vector Store Result", result)
}
return 0, nil
}
// proxyToolCallLogProvider is a proxy implementation of ToolCallLogProvider
type proxyToolCallLogProvider struct{}
// PutLog implements the ToolCallLogProvider interface
func (p *proxyToolCallLogProvider) PutLog(ctx context.Context, callID string, name string, args json.RawMessage, taskID *int64, subtaskID *int64) (int64, error) {
terminal.PrintInfo("Tool call log saved:")
terminal.PrintKeyValue("Call ID", callID)
terminal.PrintKeyValue("Name", name)
terminal.PrintKeyValue("Args", string(args))
if taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
}
if subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
}
return 0, nil
}
// UpdateLogSuccess implements the ToolCallLogProvider interface
func (p *proxyToolCallLogProvider) UpdateLogSuccess(ctx context.Context, id int64, result string, durationSeconds float64) error {
terminal.PrintInfo("Tool call log success updated:")
terminal.PrintKeyValueFormat("ID", "%d", id)
terminal.PrintKeyValue("Result", result)
terminal.PrintKeyValueFormat("Duration Seconds", "%f", durationSeconds)
return nil
}
// UpdateLogFailed implements the ToolCallLogProvider interface
func (p *proxyToolCallLogProvider) UpdateLogFailed(ctx context.Context, id int64, result string, durationSeconds float64) error {
terminal.PrintInfo("Tool call log failed updated:")
terminal.PrintKeyValueFormat("ID", "%d", id)
terminal.PrintKeyValue("Result", result)
terminal.PrintKeyValueFormat("Duration Seconds", "%f", durationSeconds)
return nil
}
// proxyKnowledgeProvider is a proxy implementation of KnowledgeProvider
type proxyKnowledgeProvider struct{}
// KnowledgeDocumentCreated implements the KnowledgeProvider interface
func (p *proxyKnowledgeProvider) KnowledgeDocumentCreated(ctx context.Context, doc *model.KnowledgeDocument) {
terminal.PrintInfo("Knowledge document created:")
terminal.PrintKeyValue("ID", doc.ID)
terminal.PrintKeyValue("Type", string(doc.DocType))
terminal.PrintKeyValue("Content", doc.Content)
terminal.PrintKeyValue("Question", doc.Question)
if doc.Description != nil {
terminal.PrintKeyValue("Description", *doc.Description)
}
if doc.FlowID != nil {
terminal.PrintKeyValueFormat("Flow ID", "%d", *doc.FlowID)
}
if doc.TaskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *doc.TaskID)
}
if doc.SubtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *doc.SubtaskID)
}
if doc.GuideType != nil {
terminal.PrintKeyValue("Guide Type", string(*doc.GuideType))
}
if doc.AnswerType != nil {
terminal.PrintKeyValue("Answer Type", string(*doc.AnswerType))
}
if doc.CodeLang != nil {
terminal.PrintKeyValue("Code Lang", *doc.CodeLang)
}
terminal.PrintKeyValueFormat("Part Size", "%d", doc.PartSize)
terminal.PrintKeyValueFormat("Total Size", "%d", doc.TotalSize)
terminal.PrintKeyValue("Manual", fmt.Sprintf("%t", doc.Manual))
terminal.PrintKeyValueFormat("User ID", "%d", doc.UserID)
}
backend/cmd/ftester/mocks/tools.go
+645
View File
@@ -0,0 +1,645 @@
package mocks
import (
"encoding/json"
"fmt"
"strings"
"pentagi/pkg/terminal"
"pentagi/pkg/tools"
)
// MockResponse generates a mock response for a function
func MockResponse(funcName string, args json.RawMessage) (string, error) {
var resultObj any
switch funcName {
case tools.TerminalToolName:
var termArgs tools.TerminalAction
if err := json.Unmarshal(args, &termArgs); err != nil {
return "", fmt.Errorf("error unmarshaling terminal arguments: %w", err)
}
terminal.PrintMock("Would execute terminal command:")
terminal.PrintKeyValue("Command", termArgs.Input)
terminal.PrintKeyValue("Working directory", termArgs.Cwd)
terminal.PrintKeyValueFormat("Timeout", "%d seconds", termArgs.Timeout.Int64())
terminal.PrintKeyValueFormat("Detach", "%v", termArgs.Detach.Bool())
if termArgs.Detach.Bool() {
resultObj = "Command executed successfully in the background mode"
} else {
resultObj = fmt.Sprintf("Mock output for command: %s\nCommand executed successfully", termArgs.Input)
}
case tools.FileToolName:
var fileArgs tools.FileAction
if err := json.Unmarshal(args, &fileArgs); err != nil {
return "", fmt.Errorf("error unmarshaling file arguments: %w", err)
}
terminal.PrintMock("File operation:")
terminal.PrintKeyValue("Operation", string(fileArgs.Action))
terminal.PrintKeyValue("Path", fileArgs.Path)
if fileArgs.Action == tools.ReadFile {
resultObj = fmt.Sprintf("Mock content of file: %s\nThis is a sample content that would be read from the file.\nIt contains multiple lines to simulate a real file.", fileArgs.Path)
} else {
resultObj = fmt.Sprintf("file %s written successfully", fileArgs.Path)
}
case tools.BrowserToolName:
var browserArgs tools.Browser
if err := json.Unmarshal(args, &browserArgs); err != nil {
return "", fmt.Errorf("error unmarshaling browser arguments: %w", err)
}
terminal.PrintMock("Browser action:")
terminal.PrintKeyValue("Action", string(browserArgs.Action))
terminal.PrintKeyValue("URL", browserArgs.Url)
switch browserArgs.Action {
case tools.Markdown:
resultObj = fmt.Sprintf("# Mock page for %s\n\n## Introduction\n\nThis is a mock page content that simulates what the real browser tool would return in markdown format.\n\n## Main Content\n\nHere is some example text that would appear on the page.\n\n* List item 1\n* List item 2\n* List item 3\n\n## Conclusion\n\nThis mock content is designed to look like real markdown content from a web page.", browserArgs.Url)
case tools.HTML:
resultObj = fmt.Sprintf("<!DOCTYPE html>\n<html>\n<head>\n <title>Mock Page for %s</title>\n</head>\n<body>\n <h1>Mock HTML Content</h1>\n <p>This is a mock HTML page that simulates what the real browser tool would return.</p>\n <ul>\n <li>HTML Element 1</li>\n <li>HTML Element 2</li>\n <li>HTML Element 3</li>\n </ul>\n</body>\n</html>", browserArgs.Url)
case tools.Links:
resultObj = fmt.Sprintf("Links list from URL '%s'\n[Homepage](https://example.com)\n[About Us](https://example.com/about)\n[Products](https://example.com/products)\n[Documentation](https://example.com/docs)\n[Contact](https://example.com/contact)", browserArgs.Url)
}
case tools.GoogleToolName:
var searchArgs tools.SearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search arguments: %w", err)
}
terminal.PrintMock("Google search:")
terminal.PrintKeyValue("Query", searchArgs.Query)
terminal.PrintKeyValueFormat("Max results", "%d", searchArgs.MaxResults.Int())
var builder strings.Builder
for i := 1; i <= min(searchArgs.MaxResults.Int(), 5); i++ {
builder.WriteString(fmt.Sprintf("# %d. Mock Google Result %d for '%s'\n\n", i, i, searchArgs.Query))
builder.WriteString(fmt.Sprintf("## URL\nhttps://example.com/result%d\n\n", i))
builder.WriteString(fmt.Sprintf("## Snippet\n\nThis is a detailed mock snippet for search result %d that matches your query '%s'. It contains relevant information that would be returned by the real Google search API.\n\n", i, searchArgs.Query))
}
resultObj = builder.String()
case tools.DuckDuckGoToolName:
var searchArgs tools.SearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search arguments: %w", err)
}
terminal.PrintMock("DuckDuckGo search:")
terminal.PrintKeyValue("Query", searchArgs.Query)
terminal.PrintKeyValueFormat("Max results", "%d", searchArgs.MaxResults.Int())
var builder strings.Builder
for i := 1; i <= min(searchArgs.MaxResults.Int(), 5); i++ {
builder.WriteString(fmt.Sprintf("# %d. Mock DuckDuckGo Result %d for '%s'\n\n", i, i, searchArgs.Query))
builder.WriteString(fmt.Sprintf("## URL\nhttps://example.com/duckduckgo/result%d\n\n", i))
builder.WriteString(fmt.Sprintf("## Description\n\nThis is a detailed mock description for search result %d that matches your query '%s'. DuckDuckGo would provide this kind of anonymous search result.\n\n", i, searchArgs.Query))
if i < min(searchArgs.MaxResults.Int(), 5) {
builder.WriteString("---\n\n")
}
}
resultObj = builder.String()
case tools.TavilyToolName:
var searchArgs tools.SearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search arguments: %w", err)
}
terminal.PrintMock("Tavily search:")
terminal.PrintKeyValue("Query", searchArgs.Query)
terminal.PrintKeyValueFormat("Max results", "%d", searchArgs.MaxResults.Int())
var builder strings.Builder
builder.WriteString("# Answer\n\n")
builder.WriteString(fmt.Sprintf("This is a comprehensive answer to your query '%s' that would be generated by Tavily AI. It synthesizes information from multiple sources to provide you with the most relevant information.\n\n", searchArgs.Query))
builder.WriteString("# Links\n\n")
for i := 1; i <= min(searchArgs.MaxResults.Int(), 3); i++ {
builder.WriteString(fmt.Sprintf("## %d. Mock Tavily Result %d\n\n", i, i))
builder.WriteString(fmt.Sprintf("* URL https://example.com/tavily/result%d\n", i))
builder.WriteString(fmt.Sprintf("* Match score %.3f\n\n", 0.95-float64(i-1)*0.1))
builder.WriteString(fmt.Sprintf("### Short content\n\nHere is a brief summary of the content from this search result related to '%s'.\n\n", searchArgs.Query))
builder.WriteString(fmt.Sprintf("### Content\n\nThis is the full detailed content that would be retrieved from the URL. It contains comprehensive information about '%s' that helps answer your query with specific facts and data points that would be relevant to your search.\n\n", searchArgs.Query))
}
resultObj = builder.String()
case tools.TraversaalToolName:
var searchArgs tools.SearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search arguments: %w", err)
}
terminal.PrintMock("Traversaal search:")
terminal.PrintKeyValue("Query", searchArgs.Query)
terminal.PrintKeyValueFormat("Max results", "%d", searchArgs.MaxResults.Int())
var builder strings.Builder
builder.WriteString("# Answer\n\n")
builder.WriteString(fmt.Sprintf("Here is the Traversaal answer to your query '%s'. Traversaal provides concise answers based on web information with relevant links for further exploration.\n\n", searchArgs.Query))
builder.WriteString("# Links\n\n")
for i := 1; i <= min(searchArgs.MaxResults.Int(), 5); i++ {
builder.WriteString(fmt.Sprintf("%d. https://example.com/traversaal/resource%d\n", i, i))
}
resultObj = builder.String()
case tools.PerplexityToolName:
var searchArgs tools.SearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search arguments: %w", err)
}
terminal.PrintMock("Perplexity search:")
terminal.PrintKeyValue("Query", searchArgs.Query)
terminal.PrintKeyValueFormat("Max results", "%d", searchArgs.MaxResults.Int())
var builder strings.Builder
builder.WriteString("# Answer\n\n")
builder.WriteString(fmt.Sprintf("This is a detailed research report from Perplexity AI about '%s'. Perplexity provides comprehensive answers by synthesizing information from various sources and augmenting it with AI analysis.\n\n", searchArgs.Query))
builder.WriteString("The query you've asked about requires examining multiple perspectives and sources. Based on recent information, here's a thorough analysis of the topic with key insights and developments.\n\n")
builder.WriteString("First, it's important to understand the background of this subject. Several authoritative sources indicate that this is an evolving area with recent developments. The most current research suggests that...\n\n")
builder.WriteString("\n\n# Citations\n\n")
for i := 1; i <= min(searchArgs.MaxResults.Int(), 3); i++ {
builder.WriteString(fmt.Sprintf("%d. https://example.com/perplexity/citation%d\n", i, i))
}
resultObj = builder.String()
case tools.SploitusToolName:
var sploitusArgs tools.SploitusAction
if err := json.Unmarshal(args, &sploitusArgs); err != nil {
return "", fmt.Errorf("error unmarshaling sploitus arguments: %w", err)
}
exploitType := sploitusArgs.ExploitType
if exploitType == "" {
exploitType = "exploits"
}
terminal.PrintMock("Sploitus search:")
terminal.PrintKeyValue("Query", sploitusArgs.Query)
terminal.PrintKeyValue("Exploit type", exploitType)
terminal.PrintKeyValue("Sort", sploitusArgs.Sort)
terminal.PrintKeyValueFormat("Max results", "%d", sploitusArgs.MaxResults.Int())
var builder strings.Builder
builder.WriteString("# Sploitus Search Results\n\n")
builder.WriteString(fmt.Sprintf("**Query:** `%s` \n", sploitusArgs.Query))
builder.WriteString(fmt.Sprintf("**Type:** %s \n", exploitType))
builder.WriteString(fmt.Sprintf("**Total matches on Sploitus:** %d\n\n", 200))
builder.WriteString("---\n\n")
maxResults := min(sploitusArgs.MaxResults.Int(), 3)
if exploitType == "tools" {
builder.WriteString(fmt.Sprintf("## Security Tools (showing up to %d)\n\n", maxResults))
for i := 1; i <= maxResults; i++ {
builder.WriteString(fmt.Sprintf("### %d. SQLMap - Automated SQL Injection Tool\n\n", i))
builder.WriteString("**URL:** https://github.com/sqlmapproject/sqlmap \n")
builder.WriteString("**Download:** https://github.com/sqlmapproject/sqlmap \n")
builder.WriteString("**Source Type:** kitploit \n")
builder.WriteString("**ID:** KITPLOIT:123456789 \n")
builder.WriteString("\n---\n\n")
}
} else {
builder.WriteString(fmt.Sprintf("## Exploits (showing up to %d)\n\n", maxResults))
builder.WriteString("### 1. SSTI-to-RCE-Python-Eval-Bypass\n\n")
builder.WriteString("**URL:** https://github.com/Rohitberiwala/SSTI-to-RCE-Python-Eval-Bypass \n")
builder.WriteString("**CVSS Score:** 5.8 \n")
builder.WriteString("**Type:** githubexploit \n")
builder.WriteString("**Published:** 2026-02-23 \n")
builder.WriteString("**ID:** 1A2B3C4D-5E6F-7G8H-9I0J-1K2L3M4N5O6P \n")
builder.WriteString("**Language:** python \n")
builder.WriteString("\n---\n\n")
if maxResults >= 2 {
builder.WriteString("### 2. Apache Struts CVE-2024-53677 RCE\n\n")
builder.WriteString("**URL:** https://github.com/example/struts-exploit \n")
builder.WriteString("**CVSS Score:** 9.8 \n")
builder.WriteString("**Type:** packetstorm \n")
builder.WriteString("**Published:** 2026-02-15 \n")
builder.WriteString("**ID:** PACKETSTORM:215999 \n")
builder.WriteString("**Language:** bash \n")
builder.WriteString("\n---\n\n")
}
if maxResults >= 3 {
builder.WriteString("### 3. Linux Kernel Privilege Escalation\n\n")
builder.WriteString("**URL:** https://www.exploit-db.com/exploits/51234 \n")
builder.WriteString("**CVSS Score:** 7.8 \n")
builder.WriteString("**Type:** metasploit \n")
builder.WriteString("**Published:** 2026-01-28 \n")
builder.WriteString("**ID:** MSF:EXPLOIT-LINUX-LOCAL-KERNEL-51234- \n")
builder.WriteString("**Language:** RUBY \n")
builder.WriteString("\n---\n\n")
}
}
resultObj = builder.String()
case tools.SearxngToolName:
var searchArgs tools.SearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search arguments: %w", err)
}
terminal.PrintMock("Searxng search:")
terminal.PrintKeyValue("Query", searchArgs.Query)
terminal.PrintKeyValueFormat("Max results", "%d", searchArgs.MaxResults.Int())
var builder strings.Builder
builder.WriteString("# Search Results\n\n")
builder.WriteString(fmt.Sprintf("This is a mock response from the Searxng meta search engine for query '%s'. In a real implementation, this would return actual search results aggregated from multiple search engines with customizable categories, language settings, and safety filters.\n\n", searchArgs.Query))
builder.WriteString("## Results\n\n")
for i := 1; i <= min(searchArgs.MaxResults.Int(), 5); i++ {
builder.WriteString(fmt.Sprintf("%d. **Mock Result %d** - Mock title about %s\n", i, i, searchArgs.Query))
builder.WriteString(fmt.Sprintf(" URL: https://example.com/searxng/result%d\n", i))
builder.WriteString(fmt.Sprintf(" Source: Mock Engine %d\n", i))
builder.WriteString(fmt.Sprintf(" Content: This is a mock content snippet that would appear in a real Searxng search result. It contains relevant information about '%s' that helps answer your query.\n\n", searchArgs.Query))
}
builder.WriteString("## Quick Answers\n\n")
builder.WriteString("- Mock answer: Based on your query, here's a quick answer that Searxng might provide.\n")
builder.WriteString("- Related search: You might also be interested in searching for related terms.\n\n")
builder.WriteString("## Related Searches\n\n")
builder.WriteString(fmt.Sprintf("- %s alternatives\n", searchArgs.Query))
builder.WriteString(fmt.Sprintf("- %s tutorial\n", searchArgs.Query))
builder.WriteString(fmt.Sprintf("- %s vs other search engines\n", searchArgs.Query))
resultObj = builder.String()
case tools.SearchToolName:
var searchArgs tools.ComplexSearch
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling complex search arguments: %w", err)
}
terminal.PrintMock("Complex search:")
terminal.PrintKeyValue("Question", searchArgs.Question)
resultObj = fmt.Sprintf("# Comprehensive Search Results for: '%s'\n\n## Summary\nThis is a comprehensive answer to your complex question based on multiple search engines and memory sources. The researcher team has compiled the most relevant information from various sources.\n\n## Key Findings\n1. Finding one: Important information related to your query\n2. Finding two: Additional context that helps answer your question\n3. Finding three: Specific details from technical documentation\n\n## Sources\n- Web search (Google, DuckDuckGo)\n- Technical documentation\n- Academic papers\n- Long-term memory results\n\n## Conclusion\nBased on all available information, here is the complete answer to your question with code examples, command samples, and specific technical details as requested.", searchArgs.Question)
case tools.SearchResultToolName:
var searchResultArgs tools.SearchResult
if err := json.Unmarshal(args, &searchResultArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search result arguments: %w", err)
}
terminal.PrintMock("Search result received:")
terminal.PrintKeyValueFormat("Content length", "%d chars", len(searchResultArgs.Result))
resultObj = map[string]any{
"status": "success",
"message": "Search results processed and delivered successfully",
}
case tools.MemoristToolName:
var memoristArgs tools.MemoristAction
if err := json.Unmarshal(args, &memoristArgs); err != nil {
return "", fmt.Errorf("error unmarshaling memorist arguments: %w", err)
}
terminal.PrintMock("Memorist question:")
terminal.PrintKeyValue("Question", memoristArgs.Question)
if memoristArgs.TaskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", memoristArgs.TaskID.Int64())
}
if memoristArgs.SubtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", memoristArgs.SubtaskID.Int64())
}
resultObj = fmt.Sprintf("# Archivist Memory Results\n\n## Question\n%s\n\n## Retrieved Information\nThe archivist has searched through all past work and tasks and found the following relevant information:\n\n1. On [date], a similar task was performed with the following approach...\n2. The team previously encountered this issue and resolved it by...\n3. Related documentation was created during project [X] that explains...\n\n## Historical Context\nThis question relates to work that was done approximately [time period] ago, and involved the following components and techniques...\n\n## Recommended Next Steps\nBased on historical information, the most effective approach would be to...", memoristArgs.Question)
case tools.MemoristResultToolName:
var memoristResultArgs tools.MemoristResult
if err := json.Unmarshal(args, &memoristResultArgs); err != nil {
return "", fmt.Errorf("error unmarshaling memorist result arguments: %w", err)
}
terminal.PrintMock("Memorist result received:")
terminal.PrintKeyValueFormat("Content length", "%d chars", len(memoristResultArgs.Result))
resultObj = map[string]any{
"status": "success",
"message": "Memory search results processed and delivered successfully",
}
case tools.SearchInMemoryToolName:
var searchMemoryArgs tools.SearchInMemoryAction
if err := json.Unmarshal(args, &searchMemoryArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search memory arguments: %w", err)
}
terminal.PrintMock("Search in memory:")
terminal.PrintKeyValueFormat("Questions count", "%d", len(searchMemoryArgs.Questions))
for i, q := range searchMemoryArgs.Questions {
terminal.PrintKeyValueFormat(fmt.Sprintf("Question %d", i+1), "%s", q)
}
if searchMemoryArgs.TaskID != nil {
terminal.PrintKeyValueFormat("Task ID filter", "%d", searchMemoryArgs.TaskID.Int64())
}
if searchMemoryArgs.SubtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID filter", "%d", searchMemoryArgs.SubtaskID.Int64())
}
questionsText := strings.Join(searchMemoryArgs.Questions, " | ")
var builder strings.Builder
builder.WriteString("# Match score 0.92\n\n")
if searchMemoryArgs.TaskID != nil {
builder.WriteString(fmt.Sprintf("# Task ID %d\n\n", searchMemoryArgs.TaskID.Int64()))
}
if searchMemoryArgs.SubtaskID != nil {
builder.WriteString(fmt.Sprintf("# Subtask ID %d\n\n", searchMemoryArgs.SubtaskID.Int64()))
}
builder.WriteString("# Tool Name 'terminal'\n\n")
builder.WriteString("# Tool Description\n\nCalls a terminal command in blocking mode. Use timeout=0 or a negative value to apply the configured server default timeout. Explicit positive values are accepted up to 10800 seconds (3 hours); values outside this range are replaced by the server default\n\n")
builder.WriteString("# Chunk\n\n")
builder.WriteString(fmt.Sprintf("This is a memory chunk related to your questions '%s'. It contains information about previous commands, outputs, and relevant context that was stored in the vector database.\n\n", questionsText))
builder.WriteString("---------------------------\n")
builder.WriteString("# Match score 0.85\n\n")
builder.WriteString("# Tool Name 'file'\n\n")
builder.WriteString("# Chunk\n\n")
builder.WriteString("This is another memory chunk that provides additional context to your questions. It contains information about file operations and relevant content changes.\n")
builder.WriteString("---------------------------\n")
resultObj = builder.String()
case tools.SearchGuideToolName:
var searchGuideArgs tools.SearchGuideAction
if err := json.Unmarshal(args, &searchGuideArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search guide arguments: %w", err)
}
terminal.PrintMock("Search guide:")
terminal.PrintKeyValueFormat("Questions count", "%d", len(searchGuideArgs.Questions))
for i, q := range searchGuideArgs.Questions {
terminal.PrintKeyValueFormat(fmt.Sprintf("Question %d", i+1), "%s", q)
}
terminal.PrintKeyValue("Guide type", searchGuideArgs.Type)
questionsText := strings.Join(searchGuideArgs.Questions, " | ")
if searchGuideArgs.Type == "pentest" {
resultObj = fmt.Sprintf("# Original Guide Type: pentest\n\n# Original Guide Questions\n\n%s\n\n## Penetration Testing Guide\n\nThis guide provides a step-by-step approach for conducting a penetration test on the target system.\n\n### 1. Reconnaissance\n- Gather information about the target using OSINT tools\n- Identify potential entry points and attack surfaces\n\n### 2. Scanning\n- Use tools like Nmap to scan for open ports and services\n- Identify vulnerabilities using automated scanners\n\n### 3. Exploitation\n- Attempt to exploit identified vulnerabilities\n- Document successful attack vectors\n\n### 4. Post-Exploitation\n- Maintain access and explore the system\n- Identify sensitive data and potential lateral movement paths\n\n### 5. Reporting\n- Document all findings with proof of concept\n- Provide remediation recommendations\n\n", questionsText)
} else if searchGuideArgs.Type == "install" {
resultObj = fmt.Sprintf("# Original Guide Type: install\n\n# Original Guide Questions\n\n%s\n\n## Installation Guide\n\n### Prerequisites\n- Operating System: Linux/macOS/Windows\n- Required dependencies: [list]\n\n### Installation Steps\n1. Download the software from the official repository\n ```bash\n git clone https://github.com/example/software.git\n ```\n\n2. Navigate to the project directory\n ```bash\n cd software\n ```\n\n3. Install dependencies\n ```bash\n npm install\n ```\n\n4. Build the project\n ```bash\n npm run build\n ```\n\n5. Verify installation\n ```bash\n npm test\n ```\n\n### Troubleshooting\n- Common issue 1: [solution]\n- Common issue 2: [solution]\n\n", questionsText)
} else {
resultObj = fmt.Sprintf("# Original Guide Type: %s\n\n# Original Guide Questions\n\n%s\n\n## Guide Content\n\nThis is a comprehensive guide for the requested type '%s'. It contains detailed instructions, best practices, and examples tailored to your specific questions.\n\n### Section 1: Getting Started\n[Detailed content would be here]\n\n### Section 2: Main Procedures\n[Step-by-step instructions would be here]\n\n### Section 3: Advanced Techniques\n[Advanced content would be here]\n\n### Section 4: Troubleshooting\n[Common issues and solutions would be here]\n\n", searchGuideArgs.Type, questionsText, searchGuideArgs.Type)
}
case tools.StoreGuideToolName:
var storeGuideArgs tools.StoreGuideAction
if err := json.Unmarshal(args, &storeGuideArgs); err != nil {
return "", fmt.Errorf("error unmarshaling store guide arguments: %w", err)
}
terminal.PrintMock("Store guide:")
terminal.PrintKeyValue("Type", storeGuideArgs.Type)
terminal.PrintKeyValueFormat("Guide length", "%d chars", len(storeGuideArgs.Guide))
terminal.PrintKeyValue("Guide question", storeGuideArgs.Question)
resultObj = "guide stored successfully"
case tools.SearchAnswerToolName:
var searchAnswerArgs tools.SearchAnswerAction
if err := json.Unmarshal(args, &searchAnswerArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search answer arguments: %w", err)
}
terminal.PrintMock("Search answer:")
terminal.PrintKeyValueFormat("Questions count", "%d", len(searchAnswerArgs.Questions))
for i, q := range searchAnswerArgs.Questions {
terminal.PrintKeyValueFormat(fmt.Sprintf("Question %d", i+1), "%s", q)
}
terminal.PrintKeyValue("Answer type", searchAnswerArgs.Type)
questionsText := strings.Join(searchAnswerArgs.Questions, " | ")
if searchAnswerArgs.Type == "vulnerability" {
resultObj = fmt.Sprintf("# Original Answer Type: vulnerability\n\n# Original Search Questions\n\n%s\n\n## Vulnerability Details\n\n### CVE-2023-12345\n\n**Severity**: High\n\n**Affected Systems**: Linux servers running Apache 2.4.x before 2.4.56\n\n**Description**:\nA buffer overflow vulnerability in Apache HTTP Server allows attackers to execute arbitrary code via a crafted request.\n\n**Exploitation**:\nAttackers can send a specially crafted HTTP request that triggers the buffer overflow, leading to remote code execution with the privileges of the web server process.\n\n**Remediation**:\n- Update Apache HTTP Server to version 2.4.56 or later\n- Apply the security patch provided by the vendor\n- Implement network filtering to block malicious requests\n\n**References**:\n- https://example.com/cve-2023-12345\n- https://example.com/apache-advisory\n", questionsText)
} else {
resultObj = fmt.Sprintf("# Original Answer Type: %s\n\n# Original Search Questions\n\n%s\n\n## Comprehensive Answer\n\nThis is a detailed answer to your questions related to the type '%s'. The answer provides comprehensive information, examples, and best practices.\n\n### Key Points\n1. First important point about your questions\n2. Second important aspect to consider\n3. Technical details relevant to your inquiry\n\n### Examples\n```\nExample code or configuration would be here\n```\n\n### Additional Resources\n- Resource 1: [description]\n- Resource 2: [description]\n\n", searchAnswerArgs.Type, questionsText, searchAnswerArgs.Type)
}
case tools.StoreAnswerToolName:
var storeAnswerArgs tools.StoreAnswerAction
if err := json.Unmarshal(args, &storeAnswerArgs); err != nil {
return "", fmt.Errorf("error unmarshaling store answer arguments: %w", err)
}
terminal.PrintMock("Store answer:")
terminal.PrintKeyValue("Type", storeAnswerArgs.Type)
terminal.PrintKeyValueFormat("Answer length", "%d chars", len(storeAnswerArgs.Answer))
terminal.PrintKeyValue("Question", storeAnswerArgs.Question)
resultObj = "answer for question stored successfully"
case tools.SearchCodeToolName:
var searchCodeArgs tools.SearchCodeAction
if err := json.Unmarshal(args, &searchCodeArgs); err != nil {
return "", fmt.Errorf("error unmarshaling search code arguments: %w", err)
}
terminal.PrintMock("Search code:")
terminal.PrintKeyValueFormat("Questions count", "%d", len(searchCodeArgs.Questions))
for i, q := range searchCodeArgs.Questions {
terminal.PrintKeyValueFormat(fmt.Sprintf("Question %d", i+1), "%s", q)
}
terminal.PrintKeyValue("Language", searchCodeArgs.Lang)
questionsText := strings.Join(searchCodeArgs.Questions, " | ")
var mockCode string
if searchCodeArgs.Lang == "python" {
mockCode = "def example_function(param1, param2='default'):\n \"\"\"This is an example Python function that demonstrates a pattern.\n \n Args:\n param1: The first parameter\n param2: The second parameter with default value\n \n Returns:\n The processed result\n \"\"\"\n result = {}\n \n # Process the parameters\n if param1 is not None:\n result['param1'] = param1\n \n # Additional processing\n if param2 != 'default':\n result['param2'] = param2\n \n return result\n\n# Example usage\nif __name__ == '__main__':\n output = example_function('test', 'custom')\n print(output)"
} else if searchCodeArgs.Lang == "javascript" || searchCodeArgs.Lang == "js" {
mockCode = "/**\n * Example JavaScript function that demonstrates a pattern\n * @param {Object} options - Configuration options\n * @param {string} options.name - The name parameter\n * @param {number} [options.count=1] - Optional count parameter\n * @returns {Object} The processed result\n */\nfunction exampleFunction(options) {\n const { name, count = 1 } = options;\n \n // Input validation\n if (!name) {\n throw new Error('Name is required');\n }\n \n // Process the data\n const result = {\n processedName: name.toUpperCase(),\n repeatedCount: Array(count).fill(name).join(', ')\n };\n \n return result;\n}\n\n// Example usage\nconst output = exampleFunction({ name: 'test', count: 3 });\nconsole.log(output);"
} else {
mockCode = fmt.Sprintf("// Example code in %s language\n// This is a mock code snippet that would be returned from the vector database\n\n// Main function definition\nfunction exampleFunction(param) {\n // Initialization\n const result = [];\n \n // Processing logic\n for (let i = 0; i < param.length; i++) {\n result.push(processItem(param[i]));\n }\n \n return result;\n}\n\n// Helper function\nfunction processItem(item) {\n return item.transform();\n}", searchCodeArgs.Lang)
}
resultObj = fmt.Sprintf("# Original Code Questions\n\n%s\n\n# Original Code Description\n\nThis code sample demonstrates the implementation pattern for handling the specific scenarios you asked about. It includes proper error handling, input validation, and follows best practices for %s.\n\n```%s\n%s\n```\n\n", questionsText, searchCodeArgs.Lang, searchCodeArgs.Lang, mockCode)
case tools.StoreCodeToolName:
var storeCodeArgs tools.StoreCodeAction
if err := json.Unmarshal(args, &storeCodeArgs); err != nil {
return "", fmt.Errorf("error unmarshaling store code arguments: %w", err)
}
terminal.PrintMock("Store code:")
terminal.PrintKeyValue("Language", storeCodeArgs.Lang)
terminal.PrintKeyValueFormat("Code length", "%d chars", len(storeCodeArgs.Code))
terminal.PrintKeyValue("Question", storeCodeArgs.Question)
terminal.PrintKeyValue("Description", storeCodeArgs.Description)
resultObj = "code sample stored successfully"
case tools.GraphitiSearchToolName:
var searchArgs tools.GraphitiSearchAction
if err := json.Unmarshal(args, &searchArgs); err != nil {
return "", fmt.Errorf("error unmarshaling graphiti search arguments: %w", err)
}
terminal.PrintMock("Graphiti Search:")
terminal.PrintKeyValue("Search Type", searchArgs.SearchType)
terminal.PrintKeyValue("Query", searchArgs.Query)
var builder strings.Builder
switch searchArgs.SearchType {
case "recent_context":
builder.WriteString("# Recent Context\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString("**Time Window:** 2025-01-19T10:00:00Z to 2025-01-19T18:00:00Z\n\n")
builder.WriteString("## Recently Discovered Entities\n\n")
builder.WriteString("1. **Target Server** (score: 0.95)\n")
builder.WriteString(" - Labels: [IP_ADDRESS, TARGET]\n")
builder.WriteString(" - Summary: Mock target server discovered during reconnaissance\n\n")
builder.WriteString("## Recent Facts\n\n")
builder.WriteString("- **Port Discovery** (score: 0.92): Target Server HAS_PORT 80 (HTTP)\n")
builder.WriteString("- **Service Identification** (score: 0.88): Port 80 RUNS_SERVICE Apache 2.4.41\n\n")
builder.WriteString("## Recent Activity\n\n")
builder.WriteString("- **pentester_agent** (score: 0.94): Executed nmap scan on target\n")
case "successful_tools":
builder.WriteString("# Successful Tools & Techniques\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString("## Successful Executions\n\n")
builder.WriteString("1. **pentester_agent** (score: 0.96)\n")
builder.WriteString(" - Description: Executed nmap scan\n")
builder.WriteString(" - Command/Output:\n```\nnmap -sV -p 80,443 192.168.1.100\n\nPORT STATE SERVICE VERSION\n80/tcp open http Apache/2.4.41\n443/tcp open https Apache/2.4.41\n```\n\n")
builder.WriteString("2. **pentester_agent** (score: 0.92)\n")
builder.WriteString(" - Description: Successful vulnerability scan\n")
builder.WriteString(" - Command/Output:\n```\nnikto -h http://192.168.1.100\n\nFound: Outdated Apache version\nFound: Accessible .git directory\n```\n\n")
case "episode_context":
builder.WriteString("# Episode Context Results\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString("## Relevant Agent Activity\n\n")
builder.WriteString("1. **pentester_agent** (relevance: 0.94)\n")
builder.WriteString(" - Time: 2025-01-19T14:30:00Z\n")
builder.WriteString(" - Description: Analyzed web application vulnerabilities\n")
builder.WriteString(" - Content:\n```\nI have completed the reconnaissance phase and identified the following:\n- Apache web server version 2.4.41 (outdated, has known vulnerabilities)\n- Exposed .git directory at /.git/\n- Directory listing enabled on /backup/\n- Potential SQL injection in login form\n\nRecommendation: Proceed with exploitation of the .git directory first.\n```\n\n")
builder.WriteString("## Mentioned Entities\n\n")
builder.WriteString("- **192.168.1.100** (relevance: 0.96): Target IP address\n")
builder.WriteString("- **Apache 2.4.41** (relevance: 0.91): Identified web server\n")
case "entity_relationships":
builder.WriteString("# Entity Relationship Search Results\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString("## Center Node: Target Server\n")
builder.WriteString("- UUID: mock-uuid-center-123\n")
builder.WriteString("- Summary: Main target system identified during reconnaissance\n\n")
builder.WriteString("## Related Facts & Relationships\n\n")
builder.WriteString("1. **Port Relationship** (distance: 0.15)\n")
builder.WriteString(" - Fact: Target Server HAS_PORT 80\n")
builder.WriteString(" - Source: mock-uuid-center-123\n")
builder.WriteString(" - Target: mock-uuid-port-80\n\n")
builder.WriteString("2. **Service Relationship** (distance: 0.25)\n")
builder.WriteString(" - Fact: Port 80 RUNS_SERVICE Apache\n")
builder.WriteString(" - Source: mock-uuid-port-80\n")
builder.WriteString(" - Target: mock-uuid-apache\n\n")
builder.WriteString("## Related Entities\n\n")
builder.WriteString("1. **HTTP Service** (distance: 0.20)\n")
builder.WriteString(" - UUID: mock-uuid-http-service\n")
builder.WriteString(" - Labels: [SERVICE, HTTP]\n")
builder.WriteString(" - Summary: Web service running on port 80\n\n")
case "temporal_window":
builder.WriteString("# Temporal Search Results\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString(fmt.Sprintf("**Time Window:** %s to %s\n\n", searchArgs.TimeStart, searchArgs.TimeEnd))
builder.WriteString("## Facts & Relationships\n\n")
builder.WriteString("1. **Vulnerability Discovery** (score: 0.93)\n")
builder.WriteString(" - Fact: Target System HAS_VULNERABILITY CVE-2021-41773\n")
builder.WriteString(" - Created: 2025-01-19T15:00:00Z\n\n")
builder.WriteString("## Entities\n\n")
builder.WriteString("1. **CVE-2021-41773** (score: 0.95)\n")
builder.WriteString(" - UUID: mock-uuid-cve\n")
builder.WriteString(" - Labels: [VULNERABILITY, CVE]\n")
builder.WriteString(" - Summary: Apache HTTP Server path traversal vulnerability\n\n")
builder.WriteString("## Agent Responses & Tool Executions\n\n")
builder.WriteString("1. **pentester_agent** (score: 0.92)\n")
builder.WriteString(" - Description: Vulnerability assessment completed\n")
builder.WriteString(" - Created: 2025-01-19T15:30:00Z\n")
builder.WriteString(" - Content:\n```\nConfirmed CVE-2021-41773 vulnerability present on target.\nSuccessfully exploited to read /etc/passwd\n```\n\n")
case "diverse_results":
builder.WriteString("# Diverse Search Results\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString("## Communities (Context Clusters)\n\n")
builder.WriteString("1. **Reconnaissance Phase** (MMR score: 0.94)\n")
builder.WriteString(" - Summary: All activities related to initial reconnaissance and scanning\n\n")
builder.WriteString("2. **Exploitation Phase** (MMR score: 0.88)\n")
builder.WriteString(" - Summary: Activities related to vulnerability exploitation\n\n")
builder.WriteString("## Diverse Facts\n\n")
builder.WriteString("1. **Network Discovery** (MMR score: 0.91)\n")
builder.WriteString(" - Fact: Nmap scan revealed 5 open ports on target\n\n")
builder.WriteString("2. **Web Application Analysis** (MMR score: 0.85)\n")
builder.WriteString(" - Fact: Web app uses outdated framework with known XSS vulnerabilities\n\n")
case "entity_by_label":
builder.WriteString("# Entity Inventory Search\n\n")
builder.WriteString(fmt.Sprintf("**Query:** %s\n\n", searchArgs.Query))
builder.WriteString("## Matching Entities\n\n")
builder.WriteString("1. **SQL Injection Vulnerability** (score: 0.96)\n")
builder.WriteString(" - UUID: mock-uuid-sqli\n")
builder.WriteString(" - Labels: [VULNERABILITY, SQL_INJECTION]\n")
builder.WriteString(" - Summary: SQL injection found in login form\n\n")
builder.WriteString("2. **XSS Vulnerability** (score: 0.92)\n")
builder.WriteString(" - UUID: mock-uuid-xss\n")
builder.WriteString(" - Labels: [VULNERABILITY, XSS]\n")
builder.WriteString(" - Summary: Reflected XSS in search parameter\n\n")
builder.WriteString("## Associated Facts\n\n")
builder.WriteString("- **Exploit Success** (score: 0.94): SQL Injection was successfully exploited to dump database\n")
default:
builder.WriteString(fmt.Sprintf("# Mock Graphiti Search Results\n\nSearch type '%s' mock not fully implemented.\n", searchArgs.SearchType))
builder.WriteString(fmt.Sprintf("Query: %s\n\nThis would return relevant results from the temporal knowledge graph.", searchArgs.Query))
}
resultObj = builder.String()
default:
terminal.PrintMock("Generic mock response:")
terminal.PrintKeyValue("Function", funcName)
resultObj = map[string]any{
"status": "success",
"message": fmt.Sprintf("Mock result for function: %s", funcName),
"data": "This is a generic mock response for testing purposes",
}
}
var resultJSON string
// Handle string results directly
if strResult, ok := resultObj.(string); ok {
resultJSON = strResult
} else {
// Marshal object results
jsonBytes, err := json.Marshal(resultObj)
if err != nil {
return "", fmt.Errorf("error marshaling mock result: %w", err)
}
resultJSON = string(jsonBytes)
}
return resultJSON, nil
}
backend/cmd/ftester/worker/args.go
+351
View File
@@ -0,0 +1,351 @@
package worker
import (
"encoding/json"
"fmt"
"reflect"
"slices"
"strings"
"pentagi/pkg/tools"
)
// FunctionInfo represents information about a function and its arguments
type FunctionInfo struct {
Name string
Description string
Arguments []ArgumentInfo
}
// ArgumentInfo represents information about a function argument
type ArgumentInfo struct {
Name string
Type string
Description string
Required bool
Default any
Enum []any
}
// DescribeParams contains the parameters for the describe function
type DescribeParams struct {
Verbose bool `json:"verbose"`
}
var describeFuncInfo = FunctionInfo{
Name: "describe",
Description: "Display information about tasks and subtasks for the given flow ID with optional filtering",
Arguments: []ArgumentInfo{
{
Name: "verbose",
Type: "boolean",
Description: "Display full descriptions and results",
Required: false,
},
},
}
// GetAvailableFunctions returns all available functions with their descriptions
func GetAvailableFunctions() []FunctionInfo {
funcInfos := []FunctionInfo{}
for name, def := range tools.GetRegistryDefinitions() {
// Skip functions that are not available for user invocation
if !isToolAvailableForCall(name) {
continue
}
funcInfo := FunctionInfo{
Name: name,
Description: def.Description,
}
funcInfos = append(funcInfos, funcInfo)
}
// Add custom ftester functions
funcInfos = append(funcInfos, describeFuncInfo)
return funcInfos
}
// GetFunctionInfo returns information about a specific function
func GetFunctionInfo(funcName string) (FunctionInfo, error) {
// Check for custom ftester functions
if funcName == "describe" {
return describeFuncInfo, nil
}
definitions := tools.GetRegistryDefinitions()
def, ok := definitions[funcName]
if !ok {
return FunctionInfo{}, fmt.Errorf("function not found: %s", funcName)
}
// Check if the function is available for user invocation
if !isToolAvailableForCall(funcName) {
return FunctionInfo{}, fmt.Errorf("function not available for user invocation: %s", funcName)
}
fi := FunctionInfo{
Name: def.Name,
Description: def.Description,
Arguments: []ArgumentInfo{},
}
// Extract argument info from the schema
if def.Parameters == nil {
return fi, nil
}
// Handle the schema based on its actual type
var schemaObj map[string]any
// Check if it's already a map
if rawMap, ok := def.Parameters.(map[string]any); ok {
schemaObj = rawMap
} else {
// It might be a jsonschema.Schema or something else that needs to be marshaled
schemaBytes, err := json.Marshal(def.Parameters)
if err != nil {
return fi, fmt.Errorf("error marshaling schema: %w", err)
}
if err := json.Unmarshal(schemaBytes, &schemaObj); err != nil {
return fi, fmt.Errorf("error unmarshaling schema: %w", err)
}
}
// Now parse the properties
if properties, ok := schemaObj["properties"].(map[string]any); ok {
for propName, propInfo := range properties {
propMap, ok := propInfo.(map[string]any)
if !ok {
continue
}
argType := "string"
if typeInfo, ok := propMap["type"]; ok {
argType = fmt.Sprintf("%v", typeInfo)
}
description := ""
if descInfo, ok := propMap["description"]; ok {
description = fmt.Sprintf("%v", descInfo)
}
required := false
if requiredFields, ok := schemaObj["required"].([]any); ok {
for _, reqField := range requiredFields {
if reqField.(string) == propName {
required = true
break
}
}
}
defaultVal := ""
if defaultInfo, ok := propMap["default"]; ok {
defaultVal = fmt.Sprintf("%v", defaultInfo)
}
enumValues := []any{}
if enumInfo, ok := propMap["enum"]; ok {
enumValues = enumInfo.([]any)
}
fi.Arguments = append(fi.Arguments, ArgumentInfo{
Name: propName,
Type: argType,
Description: description,
Required: required,
Default: defaultVal,
Enum: enumValues,
})
}
slices.SortFunc(fi.Arguments, func(a, b ArgumentInfo) int {
return strings.Compare(a.Name, b.Name)
})
}
return fi, nil
}
// ParseFunctionArgs parses command-line arguments into a structured object for the function
func ParseFunctionArgs(funcName string, args []string) (any, error) {
// Handle describe function specially
if funcName == "describe" {
params := &DescribeParams{}
// Parse the command-line arguments for describe
for i := 0; i < len(args); i++ {
arg := args[i]
// Check if the arg starts with '-'
if !strings.HasPrefix(arg, "-") {
return nil, fmt.Errorf("invalid argument format (expected '-name'): %s", arg)
}
// Get the argument name without '-'
argName := strings.TrimPrefix(arg, "-")
switch argName {
case "verbose":
params.Verbose = true
default:
return nil, fmt.Errorf("unknown argument for describe: %s", argName)
}
}
return params, nil
}
// Get function info to check required arguments
funcInfo, err := GetFunctionInfo(funcName)
if err != nil {
return nil, err
}
// Create a map to store parsed args
parsedArgs := make(map[string]any)
// Parse the command-line arguments
for i := 0; i < len(args); i++ {
arg := args[i]
// Check if the arg starts with '-'
if !strings.HasPrefix(arg, "-") {
return nil, fmt.Errorf("invalid argument format (expected '-name'): %s", arg)
}
// Get the argument name without '-'
argName := strings.TrimPrefix(arg, "-")
// Find the argument info
var argInfo *ArgumentInfo
for _, ai := range funcInfo.Arguments {
if ai.Name == argName {
argInfo = &ai
break
}
}
if argInfo == nil {
return nil, fmt.Errorf("unknown argument: %s", argName)
}
// Check if there's a value for the argument
if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") {
// Next arg is the value
parsedArgs[argName] = args[i+1]
i++ // Skip the value in the next iteration
} else {
// Boolean flag (no value)
parsedArgs[argName] = true
}
}
// Check if all required arguments are provided
for _, arg := range funcInfo.Arguments {
if arg.Required {
if _, ok := parsedArgs[arg.Name]; !ok {
if arg.Name == "message" {
parsedArgs[arg.Name] = "dummy message"
continue
}
return nil, fmt.Errorf("missing required argument: %s", arg.Name)
}
}
}
// Find the appropriate struct type for the function
structType, err := getStructTypeForFunction(funcName)
if err != nil {
return nil, err
}
// Create a new instance of the struct
structValue := reflect.New(structType).Interface()
// Convert parsedArgs to JSON
jsonData, err := json.Marshal(parsedArgs)
if err != nil {
return nil, fmt.Errorf("error marshaling arguments: %w", err)
}
// Unmarshal JSON into the struct
err = json.Unmarshal(jsonData, structValue)
if err != nil {
return nil, fmt.Errorf("error unmarshaling arguments: %w", err)
}
return structValue, nil
}
// getStructTypeForFunction finds the appropriate struct type for a function
func getStructTypeForFunction(funcName string) (reflect.Type, error) {
// Map function names to struct types
typeMap := map[string]any{
tools.TerminalToolName: &tools.TerminalAction{},
tools.FileToolName: &tools.FileAction{},
tools.BrowserToolName: &tools.Browser{},
tools.GoogleToolName: &tools.SearchAction{},
tools.DuckDuckGoToolName: &tools.SearchAction{},
tools.TavilyToolName: &tools.SearchAction{},
tools.TraversaalToolName: &tools.SearchAction{},
tools.PerplexityToolName: &tools.SearchAction{},
tools.SearxngToolName: &tools.SearchAction{},
tools.SploitusToolName: &tools.SploitusAction{},
tools.MemoristToolName: &tools.MemoristAction{},
tools.SearchInMemoryToolName: &tools.SearchInMemoryAction{},
tools.SearchGuideToolName: &tools.SearchGuideAction{},
tools.StoreGuideToolName: &tools.StoreGuideAction{},
tools.SearchAnswerToolName: &tools.SearchAnswerAction{},
tools.StoreAnswerToolName: &tools.StoreAnswerAction{},
tools.SearchCodeToolName: &tools.SearchCodeAction{},
tools.StoreCodeToolName: &tools.StoreCodeAction{},
tools.GraphitiSearchToolName: &tools.GraphitiSearchAction{},
tools.SearchToolName: &tools.ComplexSearch{},
tools.MaintenanceToolName: &tools.MaintenanceAction{},
tools.CoderToolName: &tools.CoderAction{},
tools.PentesterToolName: &tools.PentesterAction{},
tools.AdviceToolName: &tools.AskAdvice{},
tools.FinalyToolName: &tools.Done{},
tools.AskUserToolName: &tools.AskUser{},
tools.SearchResultToolName: &tools.SearchResult{},
tools.MemoristResultToolName: &tools.MemoristResult{},
tools.MaintenanceResultToolName: &tools.TaskResult{},
tools.CodeResultToolName: &tools.CodeResult{},
tools.HackResultToolName: &tools.HackResult{},
tools.EnricherResultToolName: &tools.EnricherResult{},
tools.ReportResultToolName: &tools.TaskResult{},
tools.SubtaskListToolName: &tools.SubtaskList{},
}
structType, ok := typeMap[funcName]
if !ok {
return nil, fmt.Errorf("no struct type found for function: %s", funcName)
}
return reflect.TypeOf(structType).Elem(), nil
}
// IsToolAvailableForCall checks if a tool is available for call from the command line
func isToolAvailableForCall(toolName string) bool {
toolsMapping := tools.GetToolsByType()
availableTools := map[string]struct{}{}
for toolType, toolsList := range toolsMapping {
switch toolType {
case tools.NoneToolType, tools.StoreAgentResultToolType,
tools.StoreVectorDbToolType, tools.BarrierToolType:
continue
default:
for _, tool := range toolsList {
availableTools[tool] = struct{}{}
}
}
}
_, ok := availableTools[toolName]
return ok
}
backend/cmd/ftester/worker/executor.go
+470
View File
@@ -0,0 +1,470 @@
package worker
import (
"context"
"encoding/json"
"fmt"
"time"
"pentagi/cmd/ftester/mocks"
"pentagi/pkg/config"
"pentagi/pkg/database"
"pentagi/pkg/docker"
"pentagi/pkg/graphiti"
"pentagi/pkg/providers"
"pentagi/pkg/providers/embeddings"
"pentagi/pkg/terminal"
"pentagi/pkg/tools"
"github.com/sirupsen/logrus"
"github.com/vxcontrol/cloud/anonymizer"
"github.com/vxcontrol/cloud/anonymizer/patterns"
"github.com/vxcontrol/langchaingo/vectorstores/pgvector"
)
type agentTool struct {
handler tools.ExecutorHandler
}
func (at *agentTool) Handle(ctx context.Context, name string, args json.RawMessage) (string, error) {
if at.handler == nil {
return "", fmt.Errorf("handler for tool %s is not set", name)
}
return at.handler(ctx, name, args)
}
func (at *agentTool) IsAvailable() bool {
return at.handler != nil
}
// toolExecutor holds the necessary data for creating and managing tools
type toolExecutor struct {
flowExecutor tools.FlowToolsExecutor
replacer anonymizer.Replacer
cfg *config.Config
db database.Querier
dockerClient docker.DockerClient
handlers providers.FlowProviderHandlers
store *pgvector.Store
embedder embeddings.Embedder
graphitiClient *graphiti.Client
proxies mocks.ProxyProviders
userID int64
flowID int64
taskID *int64
subtaskID *int64
}
// newToolExecutor creates a new executor with the given parameters
func newToolExecutor(
flowExecutor tools.FlowToolsExecutor,
cfg *config.Config,
db database.Querier,
dockerClient docker.DockerClient,
handlers providers.FlowProviderHandlers,
proxies mocks.ProxyProviders,
userID, flowID int64,
taskID, subtaskID *int64,
embedder embeddings.Embedder,
graphitiClient *graphiti.Client,
) (*toolExecutor, error) {
var store *pgvector.Store
if embedder.IsAvailable() {
s, err := pgvector.New(
context.Background(),
pgvector.WithConnectionURL(cfg.DatabaseURL),
pgvector.WithEmbedder(embedder),
pgvector.WithCollectionName("langchain"),
)
if err != nil {
logrus.WithError(err).Error("failed to create pgvector store")
} else {
store = &s
}
}
allPatterns, err := patterns.LoadPatterns(patterns.PatternListTypeAll)
if err != nil {
return nil, fmt.Errorf("failed to load all patterns: %v", err)
}
// combine with config secret patterns
allPatterns.Patterns = append(allPatterns.Patterns, cfg.GetSecretPatterns()...)
replacer, err := anonymizer.NewReplacer(allPatterns.Regexes(), allPatterns.Names())
if err != nil {
return nil, fmt.Errorf("failed to create replacer: %v", err)
}
return &toolExecutor{
flowExecutor: flowExecutor,
replacer: replacer,
cfg: cfg,
db: db,
dockerClient: dockerClient,
handlers: handlers,
store: store,
embedder: embedder,
graphitiClient: graphitiClient,
proxies: proxies,
userID: userID,
flowID: flowID,
taskID: taskID,
subtaskID: subtaskID,
}, nil
}
// GetTool returns the appropriate tool for a given function name
func (te *toolExecutor) GetTool(ctx context.Context, funcName string) (tools.Tool, error) {
// Get primary container for terminal/file operations (only when needed)
var containerID int64
var containerLID string
requiresContainer := funcName == tools.TerminalToolName || funcName == tools.FileToolName
if requiresContainer {
cnt, err := te.db.GetFlowPrimaryContainer(ctx, te.flowID)
if err != nil {
return nil, fmt.Errorf("failed to get primary container for flow %d: %w", te.flowID, err)
}
containerID = cnt.ID
containerLID = cnt.LocalID.String
}
// Check which tool to create based on function name
switch funcName {
case tools.TerminalToolName:
return tools.NewTerminalTool(
te.flowID,
te.taskID,
te.subtaskID,
containerID,
containerLID,
te.dockerClient,
te.proxies.GetTermLogProvider(),
time.Duration(te.cfg.TerminalToolTimeout)*time.Second,
), nil
case tools.FileToolName:
// For file operations - uses the same terminal tool
return tools.NewTerminalTool(
te.flowID,
te.taskID,
te.subtaskID,
containerID,
containerLID,
te.dockerClient,
te.proxies.GetTermLogProvider(),
time.Duration(te.cfg.TerminalToolTimeout)*time.Second,
), nil
case tools.BrowserToolName:
return tools.NewBrowserTool(
te.flowID,
te.taskID,
te.subtaskID,
te.cfg.DataDir,
te.cfg.ScraperPrivateURL,
te.cfg.ScraperPublicURL,
te.proxies.GetScreenshotProvider(),
), nil
case tools.GoogleToolName:
return tools.NewGoogleTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
), nil
case tools.DuckDuckGoToolName:
return tools.NewDuckDuckGoTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
), nil
case tools.TavilyToolName:
return tools.NewTavilyTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
te.GetSummarizer(),
), nil
case tools.TraversaalToolName:
return tools.NewTraversaalTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
), nil
case tools.PerplexityToolName:
return tools.NewPerplexityTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
te.GetSummarizer(),
), nil
case tools.SearxngToolName:
return tools.NewSearxngTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
te.GetSummarizer(),
), nil
case tools.SploitusToolName:
return tools.NewSploitusTool(
te.cfg,
te.flowID,
te.taskID,
te.subtaskID,
te.proxies.GetSearchLogProvider(),
), nil
case tools.SearchInMemoryToolName:
return tools.NewMemoryTool(
te.flowID,
te.store,
te.proxies.GetVectorStoreLogProvider(),
), nil
case tools.SearchGuideToolName:
return tools.NewGuideTool(
te.userID,
te.flowID,
te.taskID,
te.subtaskID,
te.replacer,
te.store,
te.embedder,
te.db,
te.cfg.EmbeddingMaxTextBytes,
te.proxies.GetVectorStoreLogProvider(),
te.proxies.GetKnowledgeProvider(),
), nil
case tools.SearchAnswerToolName:
return tools.NewSearchTool(
te.userID,
te.flowID,
te.taskID,
te.subtaskID,
te.replacer,
te.store,
te.embedder,
te.db,
te.cfg.EmbeddingMaxTextBytes,
te.proxies.GetVectorStoreLogProvider(),
te.proxies.GetKnowledgeProvider(),
), nil
case tools.SearchCodeToolName:
return tools.NewCodeTool(
te.userID,
te.flowID,
te.taskID,
te.subtaskID,
te.replacer,
te.store,
te.embedder,
te.db,
te.cfg.EmbeddingMaxTextBytes,
te.proxies.GetVectorStoreLogProvider(),
te.proxies.GetKnowledgeProvider(),
), nil
case tools.GraphitiSearchToolName:
return tools.NewGraphitiSearchTool(
te.flowID,
te.taskID,
te.subtaskID,
te.graphitiClient,
), nil
// AI Agent tools
case tools.AdviceToolName:
var handler tools.ExecutorHandler
if te.handlers != nil {
if te.taskID != nil && te.subtaskID != nil {
var err error
handler, err = te.handlers.GetAskAdviceHandler(ctx, te.taskID, te.subtaskID)
if err != nil {
terminal.PrintWarning("Failed to get advice handler: %v", err)
}
} else {
terminal.PrintWarning("No task or subtask ID provided for advice tool")
}
}
return &agentTool{handler: handler}, nil
case tools.CoderToolName:
var handler tools.ExecutorHandler
if te.handlers != nil {
if te.taskID != nil && te.subtaskID != nil {
var err error
handler, err = te.handlers.GetCoderHandler(ctx, te.taskID, te.subtaskID)
if err != nil {
terminal.PrintWarning("Failed to get coder handler: %v", err)
}
} else {
terminal.PrintWarning("No task or subtask ID provided for coder tool")
}
}
return &agentTool{handler: handler}, nil
case tools.MaintenanceToolName:
var handler tools.ExecutorHandler
if te.handlers != nil {
if te.taskID != nil && te.subtaskID != nil {
var err error
handler, err = te.handlers.GetInstallerHandler(ctx, te.taskID, te.subtaskID)
if err != nil {
terminal.PrintWarning("Failed to get installer handler: %v", err)
}
} else {
terminal.PrintWarning("No task or subtask ID provided for installer tool")
}
}
return &agentTool{handler: handler}, nil
case tools.MemoristToolName:
var handler tools.ExecutorHandler
if te.handlers != nil {
if te.taskID != nil {
var err error
handler, err = te.handlers.GetMemoristHandler(ctx, te.taskID, te.subtaskID)
if err != nil {
terminal.PrintWarning("Failed to get memorist handler: %v", err)
}
} else {
terminal.PrintWarning("No task ID provided for memorist tool")
}
}
return &agentTool{handler: handler}, nil
case tools.PentesterToolName:
var handler tools.ExecutorHandler
if te.handlers != nil {
if te.taskID != nil && te.subtaskID != nil {
var err error
handler, err = te.handlers.GetPentesterHandler(ctx, te.taskID, te.subtaskID)
if err != nil {
terminal.PrintWarning("Failed to get pentester handler: %v", err)
}
} else {
terminal.PrintWarning("No task or subtask ID provided for pentester tool")
}
}
return &agentTool{handler: handler}, nil
case tools.SearchToolName:
var handler tools.ExecutorHandler
if te.handlers != nil {
var err error
if te.taskID != nil && te.subtaskID != nil {
// Use subtask specific searcher if both task and subtask IDs are available
handler, err = te.handlers.GetSubtaskSearcherHandler(ctx, te.taskID, te.subtaskID)
} else if te.taskID != nil {
// Use task specific searcher if only task ID is available
handler, err = te.handlers.GetTaskSearcherHandler(ctx, *te.taskID)
} else {
terminal.PrintWarning("No task or subtask ID provided for search tool")
}
if err != nil {
terminal.PrintWarning("Failed to get search handler: %v", err)
}
}
return &agentTool{handler: handler}, nil
// For the rest of the functions, return TODO error for now
default:
return nil, fmt.Errorf("TODO: tool for function %s is not implemented yet", funcName)
}
}
// ExecuteFunctionWrapper executes a function, choosing between mock or real execution
func (te *toolExecutor) ExecuteFunctionWrapper(ctx context.Context, funcName string, args json.RawMessage) (string, error) {
// If flowID = 0, use mock responses
if te.flowID == 0 {
terminal.PrintInfo("Using MOCK mode (flowID=0)")
return mocks.MockResponse(funcName, args)
}
// If flowID > 0, perform real function execution
terminal.PrintInfo("Using REAL mode (flowID>0)")
return te.ExecuteRealFunction(ctx, funcName, args)
}
// ExecuteRealFunction performs the real function using the executor
func (te *toolExecutor) ExecuteRealFunction(ctx context.Context, funcName string, args json.RawMessage) (string, error) {
// Execute the function
terminal.PrintInfo("Executing real function: %s", funcName)
// Get the appropriate tool for this function
tool, err := te.GetTool(ctx, funcName)
if err != nil {
return "", fmt.Errorf("error getting tool for function %s: %w", funcName, err)
}
// Check if the tool is available
if !tool.IsAvailable() {
return "", fmt.Errorf("tool for function %s is not available", funcName)
}
// Handle the function with the tool
return tool.Handle(ctx, funcName, args)
}
// ExecuteFunctionWithMode handles the general function call and displays the result
func (te *toolExecutor) ExecuteFunctionWithMode(ctx context.Context, funcName string, args any) error {
// Marshal arguments to JSON
argsJSON, err := json.Marshal(args)
if err != nil {
return fmt.Errorf("error marshaling arguments: %w", err)
}
// Nicely print function information
terminal.PrintHeader("Executing function: " + funcName)
terminal.PrintHeader("Arguments:")
terminal.PrintJSON(args)
// Execute the function (either in mock mode or real)
result, err := te.ExecuteFunctionWrapper(ctx, funcName, argsJSON)
if err != nil {
return fmt.Errorf("error executing function: %w", err)
}
// Nicely print the result
terminal.PrintHeader("\nResult:")
var resultObj any
if err := json.Unmarshal([]byte(result), &resultObj); err != nil {
// If the result is not JSON, check if it's markdown and render appropriately
terminal.PrintResult(result)
} else {
terminal.PrintJSON(resultObj)
}
terminal.PrintSuccess("\nExecution completed successfully.")
return nil
}
func (te *toolExecutor) GetSummarizer() tools.SummarizeHandler {
if te.handlers == nil {
return nil
}
return te.handlers.GetSummarizeResultHandler(te.taskID, te.subtaskID)
}
backend/cmd/ftester/worker/interactive.go
+184
View File
@@ -0,0 +1,184 @@
package worker
import (
"context"
"fmt"
"os"
"reflect"
"strconv"
"strings"
"pentagi/pkg/terminal"
)
// InteractiveFillArgs interactively fills in missing function arguments
func InteractiveFillArgs(ctx context.Context, funcName string, taskID, subtaskID *int64) (any, error) {
// Get function information
funcInfo, err := GetFunctionInfo(funcName)
if err != nil {
return nil, err
}
// Special handling for the describe function
if funcName == "describe" {
params := &DescribeParams{}
result, err := terminal.GetYesNoInputContext(ctx, "Enable verbose mode", os.Stdin)
if err != nil {
return nil, fmt.Errorf("input cancelled: %w", err)
}
params.Verbose = result
return params, nil
}
// Get the structure type for the function
structType, err := getStructTypeForFunction(funcName)
if err != nil {
return nil, err
}
// Create a new instance of the structure
structValue := reflect.New(structType).Interface()
// Create a map to store argument values
parsedArgs := make(map[string]any)
terminal.PrintHeader("Interactive argument input for function: " + funcName)
terminal.PrintInfo("Please enter values for the following arguments:")
fmt.Println()
// Request values for each argument
for _, arg := range funcInfo.Arguments {
description := arg.Description
if arg.Default != "" {
description += fmt.Sprintf(" (default: %v)", arg.Default)
}
if len(arg.Enum) > 0 {
description += fmt.Sprintf(" (enum: %v)", arg.Enum)
}
terminal.PrintHeader(description)
title := arg.Name
if arg.Required && arg.Name != "message" {
title += " (required)"
}
// Request value from the user
var value any
switch arg.Type {
case "boolean":
result, err := terminal.GetYesNoInputContext(ctx, title, os.Stdin)
if err != nil {
return nil, fmt.Errorf("input cancelled for '%s': %w", arg.Name, err)
}
value = result
case "integer", "number":
if arg.Name == "task_id" && taskID != nil {
terminal.PrintKeyValueFormat("Task ID", "%d", *taskID)
value = *taskID
break
}
if arg.Name == "subtask_id" && subtaskID != nil {
terminal.PrintKeyValueFormat("Subtask ID", "%d", *subtaskID)
value = *subtaskID
break
}
for {
strValue, err := terminal.InteractivePromptContext(ctx, title, os.Stdin)
if err != nil {
return nil, fmt.Errorf("input cancelled for '%s': %w", arg.Name, err)
}
if strValue == "" && !arg.Required {
break
}
intValue, err := strconv.Atoi(strValue)
if err != nil {
terminal.PrintError("Please enter a valid number")
continue
}
value = intValue
break
}
default: // string and other types
strValue, err := terminal.InteractivePromptContext(ctx, title, os.Stdin)
if err != nil {
return nil, fmt.Errorf("input cancelled for '%s': %w", arg.Name, err)
}
value = strValue
if value == "" && arg.Required && arg.Name == "message" {
value = "dummy message"
}
}
// If a value is entered, add it to the map
if value != nil {
parsedArgs[arg.Name] = value
}
}
// Check that all required arguments are provided
for _, arg := range funcInfo.Arguments {
if arg.Required {
if _, ok := parsedArgs[arg.Name]; !ok {
return nil, fmt.Errorf("missing required argument: %s", arg.Name)
}
}
}
// Convert parsedArgs to a structure
err = fillStructFromMap(structValue, parsedArgs)
if err != nil {
return nil, fmt.Errorf("error filling structure: %w", err)
}
return structValue, nil
}
// fillStructFromMap fills a structure with data from a map
func fillStructFromMap(structPtr any, data map[string]any) error {
val := reflect.ValueOf(structPtr).Elem()
for i := 0; i < val.NumField(); i++ {
field := val.Type().Field(i)
fieldName := field.Tag.Get("json")
// If the json tag is not set, use the field name
if fieldName == "" {
fieldName = field.Name
}
// Remove optional parts of the json tag
if comma := strings.Index(fieldName, ","); comma != -1 {
fieldName = fieldName[:comma]
}
if value, ok := data[fieldName]; ok {
fieldValue := val.Field(i)
if fieldValue.CanSet() {
switch fieldValue.Kind() {
case reflect.String:
fieldValue.SetString(value.(string))
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
fieldValue.SetInt(int64(value.(int)))
case reflect.Bool:
fieldValue.SetBool(value.(bool))
case reflect.Struct:
// For special types that may be in the tools package
// This is a simplified version that may require refinement
// depending on specific types
fmt.Printf("Complex structure field detected: %s\n", fieldName)
}
}
}
}
return nil
}
backend/cmd/ftester/worker/tester.go
+685
View File
@@ -0,0 +1,685 @@
package worker
import (
"context"
"encoding/json"
"fmt"
"pentagi/cmd/ftester/mocks"
"pentagi/pkg/config"
"pentagi/pkg/database"
"pentagi/pkg/docker"
obs "pentagi/pkg/observability"
"pentagi/pkg/observability/langfuse"
"pentagi/pkg/providers"
"pentagi/pkg/providers/provider"
"pentagi/pkg/templates"
"pentagi/pkg/terminal"
"pentagi/pkg/tools"
"github.com/sirupsen/logrus"
)
type Tester interface {
Execute(args []string) error
}
// tester represents the main testing utility for tools functions
type tester struct {
db database.Querier
cfg *config.Config
ctx context.Context
docker docker.DockerClient
providers providers.ProviderController
providerName provider.ProviderName
providerType provider.ProviderType
userID int64
flowID int64
taskID *int64
subtaskID *int64
provider provider.Provider
toolExecutor *toolExecutor
flowExecutor tools.FlowToolsExecutor
flowProvider providers.FlowProvider
proxies mocks.ProxyProviders
functions *tools.Functions
}
// NewTester creates a new instance of the tester with all necessary components
func NewTester(
db database.Querier,
cfg *config.Config,
ctx context.Context,
dockerClient docker.DockerClient,
providerController providers.ProviderController,
flowID, userID int64,
taskID, subtaskID *int64,
prvname provider.ProviderName,
) (Tester, error) {
// New provider by user
prv, err := providerController.GetProvider(ctx, prvname, userID)
if err != nil {
return nil, fmt.Errorf("failed to get provider: %w", err)
}
// Create empty functions definition
functions := &tools.Functions{}
// Initialize tools flowExecutor
flowExecutor, err := tools.NewFlowToolsExecutor(db, cfg, dockerClient, functions, userID, flowID)
if err != nil {
return nil, fmt.Errorf("failed to create flow tools executor: %w", err)
}
// Initialize proxy providers
proxies := mocks.NewProxyProviders()
// Set proxy providers to the executor
flowExecutor.SetScreenshotProvider(proxies.GetScreenshotProvider())
flowExecutor.SetAgentLogProvider(proxies.GetAgentLogProvider())
flowExecutor.SetMsgLogProvider(proxies.GetMsgLogProvider())
flowExecutor.SetSearchLogProvider(proxies.GetSearchLogProvider())
flowExecutor.SetTermLogProvider(proxies.GetTermLogProvider())
flowExecutor.SetVectorStoreLogProvider(proxies.GetVectorStoreLogProvider())
flowExecutor.SetToolCallLogProvider(proxies.GetToolCallLogProvider())
flowExecutor.SetKnowledgeProvider(proxies.GetKnowledgeProvider())
flowExecutor.SetGraphitiClient(providerController.GraphitiClient())
// Initialize tool executor
toolExecutor, err := newToolExecutor(
flowExecutor, cfg, db, dockerClient, nil, proxies,
userID, flowID, taskID, subtaskID, providerController.Embedder(),
providerController.GraphitiClient(),
)
if err != nil {
return nil, fmt.Errorf("failed to create tool executor: %w", err)
}
t := &tester{
db: db,
cfg: cfg,
ctx: ctx,
docker: dockerClient,
providers: providerController,
providerName: prvname,
providerType: prv.Type(),
userID: userID,
flowID: flowID,
taskID: taskID,
subtaskID: subtaskID,
provider: prv,
toolExecutor: toolExecutor,
flowExecutor: flowExecutor,
proxies: proxies,
functions: functions,
}
if err := t.initFlowProviderController(); err != nil {
return nil, fmt.Errorf("failed to initialize flow provider controller: %w", err)
}
return t, nil
}
// initFlowProviderController initializes the flow provider when flowID is set
func (t *tester) initFlowProviderController() error {
// When flowID=0, we're in mock mode and don't need real container or provider
// This allows testing tools functions without a running flow
if t.flowID == 0 {
return nil
}
flow, err := t.db.GetFlow(t.ctx, t.flowID)
if err != nil {
return fmt.Errorf("failed to get flow: %w", err)
}
container, err := t.db.GetFlowPrimaryContainer(t.ctx, flow.ID)
if err != nil {
return fmt.Errorf("failed to get flow primary container: %w", err)
}
user, err := t.db.GetUser(t.ctx, flow.UserID)
if err != nil {
return fmt.Errorf("failed to get user %d: %w", flow.UserID, err)
}
// Setup Langfuse observability to track the execution lifecycle
// This is critical for debugging and monitoring flow performance
// We use trace context to connect this execution with earlier/later runs
ctx, observation := obs.Observer.NewObservation(t.ctx,
langfuse.WithObservationTraceID(flow.TraceID.String),
langfuse.WithObservationTraceContext(
langfuse.WithTraceName(fmt.Sprintf("%d flow worker", flow.ID)),
langfuse.WithTraceUserID(user.Mail),
langfuse.WithTraceTags([]string{"controller"}),
langfuse.WithTraceSessionID(fmt.Sprintf("flow-%d", flow.ID)),
langfuse.WithTraceMetadata(langfuse.Metadata{
"flow_id": flow.ID,
"user_id": flow.UserID,
"user_email": user.Mail,
"user_name": user.Name,
"user_hash": user.Hash,
"user_role": user.RoleName,
"provider_name": flow.ModelProviderName,
"provider_type": flow.ModelProviderType,
}),
),
)
// Create a span for tracking the entire worker lifecycle
flowSpan := observation.Span(langfuse.WithSpanName("run tester flow worker"))
t.ctx, _ = flowSpan.Observation(ctx)
// Each flow has its own JSON configuration of allowed functions
// These determine what tools the AI can access during execution
functions := &tools.Functions{}
if err := json.Unmarshal(flow.Functions, functions); err != nil {
return wrapErrorEndSpan(t.ctx, flowSpan, "failed to unmarshal functions", err)
}
t.flowExecutor.SetFunctions(functions)
// Create a prompter for communicating with the AI model
// TODO: This will eventually be customized per user/flow
prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB
// The flow provider is the bridge between the AI model and the tools executor
// It determines which AI service (OpenAI, Claude, etc) will be used and how
// the instructions are formatted and interpreted
flowProvider, err := t.providers.LoadFlowProvider(
t.ctx,
t.providerName,
prompter,
t.flowExecutor,
t.flowID,
t.userID,
t.cfg.AskUser,
container.Image,
flow.Language,
flow.Title,
flow.ToolCallIDTemplate,
)
if err != nil {
return wrapErrorEndSpan(t.ctx, flowSpan, "failed to load flow provider", err)
}
// Connect the provider's image and embedding model to the executor
// This ensures we use the right container and vector DB configuration
t.flowExecutor.SetImage(flowProvider.Image())
t.flowExecutor.SetEmbedder(flowProvider.Embedder())
// Setup log capturing for later inspection and debugging
flowProvider.SetAgentLogProvider(t.proxies.GetAgentLogProvider())
flowProvider.SetMsgLogProvider(t.proxies.GetMsgLogProvider())
// Store references to complete the initialization chain
t.flowProvider = flowProvider
t.toolExecutor.handlers = flowProvider
return nil
}
// Execute processes command line arguments and runs the appropriate function
func (t *tester) Execute(args []string) error {
// If no args or first arg is '-help' or no args after flags processing, show general help
if len(args) == 0 || args[0] == "-help" || args[0] == "--help" {
return t.showGeneralHelp()
}
funcName := args[0]
if len(args) > 1 && (args[1] == "-help" || args[1] == "--help") {
// Show function-specific help
return t.showFunctionHelp(funcName)
}
var funcArgs any
var err error
// Handle the describe function
if funcName == "describe" {
// If no arguments are provided, use interactive mode
if len(args) == 1 {
terminal.PrintInfo("No arguments provided, using interactive mode")
funcArgs, err = InteractiveFillArgs(t.ctx, funcName, t.taskID, t.subtaskID)
} else {
// Parse describe function arguments
funcArgs, err = ParseFunctionArgs(funcName, args[1:])
}
if err != nil {
return fmt.Errorf("error parsing arguments: %w", err)
}
// Call the describe function
return t.executeDescribe(t.ctx, funcArgs.(*DescribeParams))
}
// Check if arguments are provided
if len(args) == 1 {
terminal.PrintInfo("No arguments provided, using interactive mode")
funcArgs, err = InteractiveFillArgs(t.ctx, funcName, t.taskID, t.subtaskID)
} else {
// Parse function arguments
funcArgs, err = ParseFunctionArgs(funcName, args[1:])
}
if err != nil {
return fmt.Errorf("error parsing arguments: %w", err)
}
// If flowID > 0 and the function requires terminal preparation, prepare it
if t.flowID > 0 && t.needsTeminalPrepare(funcName) {
terminal.PrintInfo("Preparing container for terminal operations...")
if err := t.flowExecutor.Prepare(t.ctx); err != nil {
return fmt.Errorf("failed to prepare executor: %w", err)
}
defer func() {
if err := t.flowExecutor.Release(t.ctx); err != nil {
terminal.PrintWarning("Failed to release executor: %v", err)
}
}()
}
// Execute the function with appropriate mode based on flowID
return t.toolExecutor.ExecuteFunctionWithMode(t.ctx, funcName, funcArgs)
}
// executeDescribe shows information about tasks and subtasks for the current flow
func (t *tester) executeDescribe(ctx context.Context, params *DescribeParams) error {
// If flowID is 0, show list of all flows
if t.flowID == 0 {
return t.executeDescribeFlows(ctx, params)
}
// If subtask_id is specified, only show that specific subtask
if t.subtaskID != nil {
return t.executeDescribeSubtask(ctx, params)
}
// If task_id is specified, show only that task and its subtasks
if t.taskID != nil {
return t.executeDescribeTask(ctx, params)
}
// Show flow info and all tasks and subtasks for this flow
return t.executeDescribeFlowTasks(ctx, params)
}
// executeDescribeFlows shows list of all flows in the system
func (t *tester) executeDescribeFlows(ctx context.Context, params *DescribeParams) error {
// Get all flows
flows, err := t.db.GetFlows(ctx)
if err != nil {
return fmt.Errorf("failed to get flows: %w", err)
}
if len(flows) == 0 {
terminal.PrintInfo("No flows found")
return nil
}
terminal.PrintHeader("Available Flows:")
terminal.PrintThickSeparator()
for _, flow := range flows {
// Always display basic info
terminal.PrintKeyValue("Flow ID", fmt.Sprintf("%d", flow.ID))
terminal.PrintKeyValue("Title", flow.Title)
terminal.PrintKeyValue("Status", string(flow.Status))
if flow.CreatedAt.Valid {
terminal.PrintKeyValue("Created At", flow.CreatedAt.Time.Format("2006-01-02 15:04:05"))
}
// Display additional info if verbose mode is enabled
if params.Verbose {
terminal.PrintKeyValue("Model", flow.Model)
terminal.PrintKeyValue("ProviderName", flow.ModelProviderName)
terminal.PrintKeyValue("ProviderType", string(flow.ModelProviderType))
terminal.PrintKeyValue("Language", flow.Language)
// Get user info who created this flow
if user, err := t.db.GetUser(ctx, flow.UserID); err == nil {
terminal.PrintKeyValue("User", fmt.Sprintf("%s (%s)", user.Name, user.Mail))
terminal.PrintKeyValue("User Role", user.RoleName)
}
}
terminal.PrintThickSeparator()
}
return nil
}
// executeDescribeSubtask shows information about a specific subtask
func (t *tester) executeDescribeSubtask(ctx context.Context, params *DescribeParams) error {
subtask, err := t.db.GetSubtask(ctx, *t.subtaskID)
if err != nil {
return fmt.Errorf("failed to get subtask: %w", err)
}
task, err := t.db.GetTask(ctx, subtask.TaskID)
if err != nil {
return fmt.Errorf("failed to get parent task: %w", err)
}
if task.FlowID != t.flowID {
return fmt.Errorf("subtask %d does not belong to flow %d", *t.subtaskID, t.flowID)
}
// Get flow information
flow, err := t.db.GetFlow(ctx, t.flowID)
if err != nil {
return fmt.Errorf("failed to get flow information: %w", err)
}
// Display flow info
terminal.PrintHeader("Flow Information")
terminal.PrintKeyValue("Flow ID", fmt.Sprintf("%d", flow.ID))
terminal.PrintKeyValue("Title", flow.Title)
terminal.PrintKeyValue("Status", string(flow.Status))
fmt.Println()
// Display task info
terminal.PrintHeader("Task Information")
terminal.PrintKeyValue("Task ID", fmt.Sprintf("%d", task.ID))
terminal.PrintKeyValue("Task Title", task.Title)
terminal.PrintKeyValue("Task Status", string(task.Status))
if params.Verbose {
terminal.PrintThinSeparator()
terminal.PrintHeader("Task Input")
terminal.RenderMarkdown(task.Input)
terminal.PrintThinSeparator()
terminal.PrintHeader("Task Result")
terminal.RenderMarkdown(task.Result)
}
fmt.Println()
// Print subtask details
terminal.PrintHeader("Subtask Information")
terminal.PrintKeyValue("Subtask ID", fmt.Sprintf("%d", subtask.ID))
terminal.PrintKeyValue("Subtask Title", subtask.Title)
terminal.PrintKeyValue("Subtask Status", string(subtask.Status))
if params.Verbose {
terminal.PrintThinSeparator()
terminal.PrintHeader("Subtask Description")
terminal.RenderMarkdown(subtask.Description)
terminal.PrintThinSeparator()
terminal.PrintHeader("Subtask Result")
terminal.RenderMarkdown(subtask.Result)
}
return nil
}
// executeDescribeTask shows information about a specific task and its subtasks
func (t *tester) executeDescribeTask(ctx context.Context, params *DescribeParams) error {
task, err := t.db.GetFlowTask(ctx, database.GetFlowTaskParams{
ID: *t.taskID,
FlowID: t.flowID,
})
if err != nil {
return fmt.Errorf("failed to get task: %w", err)
}
// Get flow information
flow, err := t.db.GetFlow(ctx, t.flowID)
if err != nil {
return fmt.Errorf("failed to get flow information: %w", err)
}
// Display flow info
terminal.PrintHeader("Flow Information")
terminal.PrintKeyValue("Flow ID", fmt.Sprintf("%d", flow.ID))
terminal.PrintKeyValue("Title", flow.Title)
terminal.PrintKeyValue("Status", string(flow.Status))
fmt.Println()
// Display task info
terminal.PrintHeader("Task Information")
terminal.PrintKeyValue("Task ID", fmt.Sprintf("%d", task.ID))
terminal.PrintKeyValue("Task Title", task.Title)
terminal.PrintKeyValue("Task Status", string(task.Status))
if params.Verbose {
terminal.PrintThinSeparator()
terminal.PrintHeader("Task Input")
terminal.RenderMarkdown(task.Input)
terminal.PrintThinSeparator()
terminal.PrintHeader("Task Result")
terminal.RenderMarkdown(task.Result)
}
fmt.Println()
// Get subtasks for this task
subtasks, err := t.db.GetFlowTaskSubtasks(ctx, database.GetFlowTaskSubtasksParams{
FlowID: t.flowID,
TaskID: *t.taskID,
})
if err != nil {
return fmt.Errorf("failed to get subtasks: %w", err)
}
if len(subtasks) == 0 {
terminal.PrintInfo("No subtasks found for this task")
return nil
}
terminal.PrintHeader(fmt.Sprintf("Subtasks for Task %d:", task.ID))
terminal.PrintThinSeparator()
for _, subtask := range subtasks {
terminal.PrintKeyValue("Subtask ID", fmt.Sprintf("%d", subtask.ID))
terminal.PrintKeyValue("Subtask Title", subtask.Title)
terminal.PrintKeyValue("Subtask Status", string(subtask.Status))
if params.Verbose {
terminal.PrintThinSeparator()
terminal.PrintHeader("Subtask Description")
terminal.RenderMarkdown(subtask.Description)
terminal.PrintThinSeparator()
terminal.PrintHeader("Subtask Result")
terminal.RenderMarkdown(subtask.Result)
}
terminal.PrintThinSeparator()
}
return nil
}
// executeDescribeFlowTasks shows information about a flow and all its tasks and subtasks
func (t *tester) executeDescribeFlowTasks(ctx context.Context, params *DescribeParams) error {
// Get flow information
flow, err := t.db.GetFlow(ctx, t.flowID)
if err != nil {
return fmt.Errorf("failed to get flow information: %w", err)
}
terminal.PrintHeader("Flow Information")
terminal.PrintKeyValue("Flow ID", fmt.Sprintf("%d", flow.ID))
terminal.PrintKeyValue("Title", flow.Title)
terminal.PrintKeyValue("Status", string(flow.Status))
terminal.PrintKeyValue("Language", flow.Language)
terminal.PrintKeyValue("Model", fmt.Sprintf("%s (%s)", flow.Model, flow.ModelProviderName))
if flow.CreatedAt.Valid {
terminal.PrintKeyValue("Created At", flow.CreatedAt.Time.Format("2006-01-02 15:04:05"))
}
fmt.Println()
// Show all tasks and subtasks for this flow
tasks, err := t.db.GetFlowTasks(ctx, t.flowID)
if err != nil {
return fmt.Errorf("failed to get tasks: %w", err)
}
if len(tasks) == 0 {
terminal.PrintInfo("No tasks found for this flow")
return nil
}
terminal.PrintHeader(fmt.Sprintf("Tasks for Flow %d:", t.flowID))
terminal.PrintThickSeparator()
for _, task := range tasks {
terminal.PrintKeyValue("Task ID", fmt.Sprintf("%d", task.ID))
terminal.PrintKeyValue("Task Title", task.Title)
terminal.PrintKeyValue("Task Status", string(task.Status))
if params.Verbose {
terminal.PrintThinSeparator()
terminal.PrintHeader("Task Input")
terminal.RenderMarkdown(task.Input)
terminal.PrintThinSeparator()
terminal.PrintHeader("Task Result")
terminal.RenderMarkdown(task.Result)
}
fmt.Println()
// Get subtasks for this task
subtasks, err := t.db.GetTaskSubtasks(ctx, task.ID)
if err != nil {
return fmt.Errorf("failed to get subtasks for task %d: %w", task.ID, err)
}
if len(subtasks) > 0 {
terminal.PrintHeader(fmt.Sprintf("Subtasks for Task %d:", task.ID))
terminal.PrintThinSeparator()
for _, subtask := range subtasks {
terminal.PrintKeyValue("Subtask ID", fmt.Sprintf("%d", subtask.ID))
terminal.PrintKeyValue("Subtask Title", subtask.Title)
terminal.PrintKeyValue("Subtask Status", string(subtask.Status))
if params.Verbose {
terminal.PrintThinSeparator()
terminal.PrintHeader("Subtask Description")
terminal.RenderMarkdown(subtask.Description)
terminal.PrintThinSeparator()
terminal.PrintHeader("Subtask Result")
terminal.RenderMarkdown(subtask.Result)
}
terminal.PrintThinSeparator()
}
} else {
terminal.PrintInfo(fmt.Sprintf("No subtasks found for Task %d", task.ID))
}
terminal.PrintThickSeparator()
}
return nil
}
// showGeneralHelp displays the general help message with a list of available functions
func (t *tester) showGeneralHelp() error {
functions := GetAvailableFunctions()
toolsByType := tools.GetToolsByType()
terminal.PrintHeader("Usage: ftester FUNCTION [ARGUMENTS]")
fmt.Println()
terminal.PrintHeader("Built-in functions:")
terminal.PrintValueFormat(" %-20s", "describe")
fmt.Printf(" - %s\n", describeFuncInfo.Description)
// Define type names for better readability
typeNames := map[tools.ToolType]string{
tools.EnvironmentToolType: "Work with terminal and files (work with environment)",
tools.SearchNetworkToolType: "Search in the internet",
tools.SearchVectorDbToolType: "Search in the Vector DB",
tools.AgentToolType: "Agents",
}
// Process each type in the order we want to display them
for _, toolType := range []tools.ToolType{
tools.SearchNetworkToolType,
tools.EnvironmentToolType,
tools.SearchVectorDbToolType,
tools.AgentToolType,
} {
// Get type name
typeName, ok := typeNames[toolType]
if !ok {
continue
}
// Get tools for this type
toolsOfType := toolsByType[toolType]
if len(toolsOfType) == 0 {
continue
}
// Print section header
fmt.Println()
terminal.PrintHeader(typeName + ":")
// Print each function in this group
for _, tool := range toolsOfType {
// Skip functions that are not available for user invocation
if !isToolAvailableForCall(tool) {
continue
}
// Find function info
var description string
for _, fn := range functions {
if fn.Name == tool {
description = fn.Description
break
}
}
terminal.PrintValueFormat(" %-20s", tool)
fmt.Printf(" - %s\n", description)
}
}
fmt.Println()
terminal.PrintInfo("For help on a specific function, use: ftester FUNCTION -help")
terminal.PrintKeyValue("Current mode", t.getModeDescription())
return nil
}
// getModeDescription returns a description of the current mode based on flowID
func (t *tester) getModeDescription() string {
if t.flowID == 0 {
return "MOCK (flowID=0)"
}
return fmt.Sprintf("REAL (flowID=%d)", t.flowID)
}
// showFunctionHelp displays help for a specific function, including its arguments
func (t *tester) showFunctionHelp(funcName string) error {
// Get function info
fnInfo, err := GetFunctionInfo(funcName)
if err != nil {
return err
}
terminal.PrintHeader(fmt.Sprintf("Function: %s", fnInfo.Name))
terminal.PrintKeyValue("Description", fnInfo.Description)
fmt.Println()
terminal.PrintHeader("Arguments:")
for _, arg := range fnInfo.Arguments {
requiredStr := ""
if arg.Required {
requiredStr = " (required)"
}
terminal.PrintValueFormat(" -%-20s", arg.Name)
fmt.Printf(" %s%s\n", arg.Description, requiredStr)
}
return nil
}
// needsTeminalPrepare determines if a function needs terminal preparation
func (t *tester) needsTeminalPrepare(funcName string) bool {
// These functions require terminal preparation
terminalFunctions := map[string]bool{
tools.TerminalToolName: true,
tools.FileToolName: true,
}
// For all other functions, no preparation is needed instead of terminal or agents functions
return terminalFunctions[funcName] || tools.GetToolTypeMapping()[funcName] == tools.AgentToolType
}
// wrapErrorEndSpan wraps an error with an end span in langfuse
func wrapErrorEndSpan(ctx context.Context, span langfuse.Span, msg string, err error) error {
logrus.WithContext(ctx).WithError(err).Error(msg)
err = fmt.Errorf("%s: %w", msg, err)
span.End(
langfuse.WithSpanStatus(err.Error()),
langfuse.WithSpanLevel(langfuse.ObservationLevelError),
)
return err
}
backend/cmd/installer/checker/checker.go
+690
View File
@@ -0,0 +1,690 @@
package checker
import (
"context"
"errors"
"fmt"
"path/filepath"
"runtime"
"sync"
"pentagi/cmd/installer/state"
"pentagi/pkg/version"
"github.com/docker/docker/client"
)
var (
InstallerVersion = version.GetBinaryVersion()
UserAgent = "PentAGI-Installer/" + InstallerVersion
)
const (
DockerComposeFile = "docker-compose.yml"
GraphitiComposeFile = "docker-compose-graphiti.yml"
LangfuseComposeFile = "docker-compose-langfuse.yml"
ObservabilityComposeFile = "docker-compose-observability.yml"
ExampleCustomConfigLLMFile = "example.custom.provider.yml"
ExampleOllamaConfigLLMFile = "example.ollama.provider.yml"
PentagiScriptFile = "/usr/local/bin/pentagi"
PentagiContainerName = "pentagi"
GraphitiContainerName = "graphiti"
Neo4jContainerName = "neo4j"
LangfuseWorkerContainerName = "langfuse-worker"
LangfuseWebContainerName = "langfuse-web"
GrafanaContainerName = "grafana"
OpenTelemetryContainerName = "otel"
DefaultImage = "debian:latest"
DefaultImageForPentest = "vxcontrol/kali-linux"
DefaultGraphitiEndpoint = "http://graphiti:8000"
DefaultLangfuseEndpoint = "http://langfuse-web:3000"
DefaultObservabilityEndpoint = "otelcol:8148"
DefaultLangfuseOtelEndpoint = "http://otelcol:4318"
DefaultUpdateServerEndpoint = "https://update.pentagi.com"
UpdatesCheckEndpoint = "/api/v1/updates/check"
MinFreeMemGB = 0.5
MinFreeMemGBForPentagi = 0.5
MinFreeMemGBForGraphiti = 2.0
MinFreeMemGBForLangfuse = 1.5
MinFreeMemGBForObservability = 1.5
MinFreeDiskGB = 5.0
MinFreeDiskGBForComponents = 10.0
MinFreeDiskGBPerComponents = 2.0
MinFreeDiskGBForWorkerImages = 25.0
)
var (
ErrAppStateNotInitialized = errors.New("appState not initialized")
ErrHandlerNotInitialized = errors.New("handler not initialized")
)
type CheckResult struct {
EnvFileExists bool `json:"env_file_exists" yaml:"env_file_exists"`
DockerApiAccessible bool `json:"docker_api_accessible" yaml:"docker_api_accessible"`
WorkerEnvApiAccessible bool `json:"worker_env_api_accessible" yaml:"worker_env_api_accessible"`
WorkerImageExists bool `json:"worker_image_exists" yaml:"worker_image_exists"`
DockerInstalled bool `json:"docker_installed" yaml:"docker_installed"`
DockerComposeInstalled bool `json:"docker_compose_installed" yaml:"docker_compose_installed"`
DockerVersion string `json:"docker_version" yaml:"docker_version"`
DockerVersionOK bool `json:"docker_version_ok" yaml:"docker_version_ok"`
DockerComposeVersion string `json:"docker_compose_version" yaml:"docker_compose_version"`
DockerComposeVersionOK bool `json:"docker_compose_version_ok" yaml:"docker_compose_version_ok"`
PentagiScriptInstalled bool `json:"pentagi_script_installed" yaml:"pentagi_script_installed"`
PentagiExtracted bool `json:"pentagi_extracted" yaml:"pentagi_extracted"`
PentagiInstalled bool `json:"pentagi_installed" yaml:"pentagi_installed"`
PentagiRunning bool `json:"pentagi_running" yaml:"pentagi_running"`
PentagiVolumesExist bool `json:"pentagi_volumes_exist" yaml:"pentagi_volumes_exist"`
GraphitiConnected bool `json:"graphiti_connected" yaml:"graphiti_connected"`
GraphitiExternal bool `json:"graphiti_external" yaml:"graphiti_external"`
GraphitiExtracted bool `json:"graphiti_extracted" yaml:"graphiti_extracted"`
GraphitiInstalled bool `json:"graphiti_installed" yaml:"graphiti_installed"`
GraphitiRunning bool `json:"graphiti_running" yaml:"graphiti_running"`
GraphitiVolumesExist bool `json:"graphiti_volumes_exist" yaml:"graphiti_volumes_exist"`
LangfuseConnected bool `json:"langfuse_connected" yaml:"langfuse_connected"`
LangfuseExternal bool `json:"langfuse_external" yaml:"langfuse_external"`
LangfuseExtracted bool `json:"langfuse_extracted" yaml:"langfuse_extracted"`
LangfuseInstalled bool `json:"langfuse_installed" yaml:"langfuse_installed"`
LangfuseRunning bool `json:"langfuse_running" yaml:"langfuse_running"`
LangfuseVolumesExist bool `json:"langfuse_volumes_exist" yaml:"langfuse_volumes_exist"`
ObservabilityConnected bool `json:"observability_connected" yaml:"observability_connected"`
ObservabilityExternal bool `json:"observability_external" yaml:"observability_external"`
ObservabilityExtracted bool `json:"observability_extracted" yaml:"observability_extracted"`
ObservabilityInstalled bool `json:"observability_installed" yaml:"observability_installed"`
ObservabilityRunning bool `json:"observability_running" yaml:"observability_running"`
SysNetworkOK bool `json:"sys_network_ok" yaml:"sys_network_ok"`
SysCPUOK bool `json:"sys_cpu_ok" yaml:"sys_cpu_ok"`
SysMemoryOK bool `json:"sys_memory_ok" yaml:"sys_memory_ok"`
SysDiskFreeSpaceOK bool `json:"sys_disk_free_space_ok" yaml:"sys_disk_free_space_ok"`
UpdateServerAccessible bool `json:"update_server_accessible" yaml:"update_server_accessible"`
InstallerIsUpToDate bool `json:"installer_is_up_to_date" yaml:"installer_is_up_to_date"`
PentagiIsUpToDate bool `json:"pentagi_is_up_to_date" yaml:"pentagi_is_up_to_date"`
GraphitiIsUpToDate bool `json:"graphiti_is_up_to_date" yaml:"graphiti_is_up_to_date"`
LangfuseIsUpToDate bool `json:"langfuse_is_up_to_date" yaml:"langfuse_is_up_to_date"`
ObservabilityIsUpToDate bool `json:"observability_is_up_to_date" yaml:"observability_is_up_to_date"`
WorkerIsUpToDate bool `json:"worker_is_up_to_date" yaml:"worker_is_up_to_date"`
// System resource details for UI display
SysCPUCount int `json:"sys_cpu_count" yaml:"sys_cpu_count"`
SysMemoryRequired float64 `json:"sys_memory_required_gb" yaml:"sys_memory_required_gb"`
SysMemoryAvailable float64 `json:"sys_memory_available_gb" yaml:"sys_memory_available_gb"`
SysDiskRequired float64 `json:"sys_disk_required_gb" yaml:"sys_disk_required_gb"`
SysDiskAvailable float64 `json:"sys_disk_available_gb" yaml:"sys_disk_available_gb"`
SysNetworkFailures []string `json:"sys_network_failures" yaml:"sys_network_failures"`
DockerErrorType DockerErrorType `json:"docker_error_type" yaml:"docker_error_type"`
EnvDirWritable bool `json:"env_dir_writable" yaml:"env_dir_writable"`
// handler controls how information is gathered. If nil, skip gathering
handler CheckHandler
}
// CheckHandler defines how to gather information into a CheckResult
type CheckHandler interface {
GatherAllInfo(ctx context.Context, c *CheckResult) error
GatherDockerInfo(ctx context.Context, c *CheckResult) error
GatherWorkerInfo(ctx context.Context, c *CheckResult) error
GatherPentagiInfo(ctx context.Context, c *CheckResult) error
GatherGraphitiInfo(ctx context.Context, c *CheckResult) error
GatherLangfuseInfo(ctx context.Context, c *CheckResult) error
GatherObservabilityInfo(ctx context.Context, c *CheckResult) error
GatherSystemInfo(ctx context.Context, c *CheckResult) error
GatherUpdatesInfo(ctx context.Context, c *CheckResult) error
}
// Delegating methods that preserve public API
func (c *CheckResult) GatherAllInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherAllInfo(ctx, c)
}
func (c *CheckResult) GatherDockerInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherDockerInfo(ctx, c)
}
func (c *CheckResult) GatherWorkerInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherWorkerInfo(ctx, c)
}
func (c *CheckResult) GatherPentagiInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherPentagiInfo(ctx, c)
}
func (c *CheckResult) GatherGraphitiInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherGraphitiInfo(ctx, c)
}
func (c *CheckResult) GatherLangfuseInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherLangfuseInfo(ctx, c)
}
func (c *CheckResult) GatherObservabilityInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherObservabilityInfo(ctx, c)
}
func (c *CheckResult) GatherSystemInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherSystemInfo(ctx, c)
}
func (c *CheckResult) GatherUpdatesInfo(ctx context.Context) error {
if c.handler == nil {
return ErrHandlerNotInitialized
}
return c.handler.GatherUpdatesInfo(ctx, c)
}
func (c *CheckResult) IsReadyToContinue() bool {
return c.EnvFileExists &&
c.EnvDirWritable &&
c.DockerApiAccessible &&
c.WorkerEnvApiAccessible &&
c.DockerComposeInstalled &&
c.DockerVersionOK &&
c.DockerComposeVersionOK &&
c.SysNetworkOK &&
c.SysCPUOK &&
c.SysMemoryOK &&
c.SysDiskFreeSpaceOK
}
// availability helpers for installer operations
// these functions centralize complex visibility/availability logic for UI
// CanStartAll returns true when at least one embedded stack is installed and not running
func (c *CheckResult) CanStartAll() bool {
if c.PentagiInstalled && !c.PentagiRunning {
return true
}
if c.GraphitiConnected && !c.GraphitiExternal && c.GraphitiInstalled && !c.GraphitiRunning {
return true
}
if c.LangfuseConnected && !c.LangfuseExternal && c.LangfuseInstalled && !c.LangfuseRunning {
return true
}
if c.ObservabilityConnected && !c.ObservabilityExternal && c.ObservabilityInstalled && !c.ObservabilityRunning {
return true
}
return false
}
// CanStopAll returns true when any compose stack is running
func (c *CheckResult) CanStopAll() bool {
return c.PentagiRunning || c.GraphitiRunning || c.LangfuseRunning || c.ObservabilityRunning
}
// CanRestartAll mirrors stop logic (requires running services)
func (c *CheckResult) CanRestartAll() bool { return c.CanStopAll() }
// CanDownloadWorker returns true when worker image is missing
func (c *CheckResult) CanDownloadWorker() bool { return !c.WorkerImageExists }
// CanUpdateWorker returns true when worker image exists but is not up to date
func (c *CheckResult) CanUpdateWorker() bool { return c.WorkerImageExists && !c.WorkerIsUpToDate }
// CanUpdateAll returns true when any installed stack has updates available
func (c *CheckResult) CanUpdateAll() bool {
if c.PentagiInstalled && !c.PentagiIsUpToDate {
return true
}
if c.GraphitiInstalled && !c.GraphitiIsUpToDate {
return true
}
if c.LangfuseInstalled && !c.LangfuseIsUpToDate {
return true
}
if c.ObservabilityInstalled && !c.ObservabilityIsUpToDate {
return true
}
return false
}
// CanUpdateInstaller returns true when installer update is available and update server accessible
func (c *CheckResult) CanUpdateInstaller() bool {
return !c.InstallerIsUpToDate && c.UpdateServerAccessible
}
// CanFactoryReset returns true when any compose stack is installed
func (c *CheckResult) CanFactoryReset() bool {
return c.PentagiInstalled || c.GraphitiInstalled || c.LangfuseInstalled || c.ObservabilityInstalled
}
// CanRemoveAll returns true when any compose stack is installed
func (c *CheckResult) CanRemoveAll() bool { return c.CanFactoryReset() }
// CanPurgeAll returns true when any compose stack is installed
func (c *CheckResult) CanPurgeAll() bool { return c.CanFactoryReset() }
// CanResetPassword returns true when PentAGI is running
func (c *CheckResult) CanResetPassword() bool { return c.PentagiRunning }
// CanInstallAll returns true when main stack is not installed yet
func (c *CheckResult) CanInstallAll() bool { return !c.PentagiInstalled }
// defaultCheckHandler provides the existing implementation of gathering logic
type defaultCheckHandler struct {
mx *sync.Mutex
appState state.State
dockerClient *client.Client
workerClient *client.Client
}
func (h *defaultCheckHandler) GatherAllInfo(ctx context.Context, c *CheckResult) error {
envPath := h.appState.GetEnvPath()
c.EnvFileExists = checkFileExists(envPath) && checkFileIsReadable(envPath)
if !c.EnvFileExists {
return fmt.Errorf("environment file %s does not exist or is not readable", envPath)
}
// check write permissions to .env directory
envDir := filepath.Dir(envPath)
c.EnvDirWritable = checkDirIsWritable(envDir)
if err := h.GatherDockerInfo(ctx, c); err != nil {
return err
}
if err := h.GatherWorkerInfo(ctx, c); err != nil {
return err
}
if err := h.GatherPentagiInfo(ctx, c); err != nil {
return err
}
if err := h.GatherGraphitiInfo(ctx, c); err != nil {
return err
}
if err := h.GatherLangfuseInfo(ctx, c); err != nil {
return err
}
if err := h.GatherObservabilityInfo(ctx, c); err != nil {
return err
}
if err := h.GatherSystemInfo(ctx, c); err != nil {
return err
}
if err := h.GatherUpdatesInfo(ctx, c); err != nil {
return err
}
return nil
}
func (h *defaultCheckHandler) GatherDockerInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
var cli *client.Client
if cli, c.DockerErrorType = createDockerClientFromEnv(ctx); c.DockerErrorType != DockerErrorNone {
c.DockerApiAccessible = false
c.DockerInstalled = c.DockerErrorType != DockerErrorNotInstalled
if c.DockerInstalled {
version := checkDockerCliVersion()
c.DockerVersion = version.Version
c.DockerVersionOK = version.Valid
}
} else {
h.dockerClient = cli
c.DockerApiAccessible = true
c.DockerInstalled = true
version := checkDockerVersion(ctx, cli)
c.DockerVersion = version.Version
c.DockerVersionOK = version.Valid
}
composeVersion := checkDockerComposeVersion()
c.DockerComposeInstalled = composeVersion.Version != ""
c.DockerComposeVersion = composeVersion.Version
c.DockerComposeVersionOK = composeVersion.Valid
return nil
}
func (h *defaultCheckHandler) GatherWorkerInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
dockerHost := getEnvVar(h.appState, "DOCKER_HOST", "")
dockerCertPath := getEnvVar(h.appState, "PENTAGI_DOCKER_CERT_PATH", "")
dockerTLSVerify := getEnvVar(h.appState, "DOCKER_TLS_VERIFY", "") != ""
cli, err := createDockerClient(dockerHost, dockerCertPath, dockerTLSVerify)
if err != nil {
// fallback to DOCKER_CERT_PATH for backward compatibility
// this handles cases where migration failed or user manually edited .env
// note: after migration, DOCKER_CERT_PATH contains container path, not host path
dockerCertPath = getEnvVar(h.appState, "DOCKER_CERT_PATH", "")
cli, err = createDockerClient(dockerHost, dockerCertPath, dockerTLSVerify)
if err != nil {
c.WorkerEnvApiAccessible = false
c.WorkerImageExists = false
return nil
}
}
h.workerClient = cli
c.WorkerEnvApiAccessible = true
pentestImage := getEnvVar(h.appState, "DOCKER_DEFAULT_IMAGE_FOR_PENTEST", DefaultImageForPentest)
c.WorkerImageExists = checkImageExists(ctx, cli, pentestImage)
return nil
}
func (h *defaultCheckHandler) GatherPentagiInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
envDir := filepath.Dir(h.appState.GetEnvPath())
dockerComposeFile := filepath.Join(envDir, DockerComposeFile)
c.PentagiExtracted = checkFileExists(dockerComposeFile) &&
checkFileExists(ExampleCustomConfigLLMFile) &&
checkFileExists(ExampleOllamaConfigLLMFile)
c.PentagiScriptInstalled = checkFileExists(PentagiScriptFile)
if h.dockerClient != nil {
exists, running := checkContainerExists(ctx, h.dockerClient, PentagiContainerName)
c.PentagiInstalled = exists
c.PentagiRunning = running
// check if pentagi-related volumes exist (indicates previous installation)
pentagiVolumes := []string{"pentagi-postgres-data", "pentagi-data", "pentagi-ssl", "scraper-ssl"}
c.PentagiVolumesExist = checkVolumesExist(ctx, h.dockerClient, pentagiVolumes)
}
return nil
}
func (h *defaultCheckHandler) GatherGraphitiInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
graphitiEnabled := getEnvVar(h.appState, "GRAPHITI_ENABLED", "")
graphitiURL := getEnvVar(h.appState, "GRAPHITI_URL", "")
c.GraphitiConnected = graphitiEnabled == "true" && graphitiURL != ""
c.GraphitiExternal = graphitiURL != DefaultGraphitiEndpoint
envDir := filepath.Dir(h.appState.GetEnvPath())
graphitiComposeFile := filepath.Join(envDir, GraphitiComposeFile)
c.GraphitiExtracted = checkFileExists(graphitiComposeFile)
if h.dockerClient != nil {
graphitiExists, graphitiRunning := checkContainerExists(ctx, h.dockerClient, GraphitiContainerName)
neo4jExists, neo4jRunning := checkContainerExists(ctx, h.dockerClient, Neo4jContainerName)
c.GraphitiInstalled = graphitiExists && neo4jExists
c.GraphitiRunning = graphitiRunning && neo4jRunning
// check if graphiti-related volumes exist (indicates previous installation)
graphitiVolumes := []string{"neo4j_data"}
c.GraphitiVolumesExist = checkVolumesExist(ctx, h.dockerClient, graphitiVolumes)
}
return nil
}
func (h *defaultCheckHandler) GatherLangfuseInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
baseURL := getEnvVar(h.appState, "LANGFUSE_BASE_URL", "")
projectID := getEnvVar(h.appState, "LANGFUSE_PROJECT_ID", "")
publicKey := getEnvVar(h.appState, "LANGFUSE_PUBLIC_KEY", "")
secretKey := getEnvVar(h.appState, "LANGFUSE_SECRET_KEY", "")
c.LangfuseConnected = baseURL != "" && projectID != "" && publicKey != "" && secretKey != ""
c.LangfuseExternal = baseURL != DefaultLangfuseEndpoint
envDir := filepath.Dir(h.appState.GetEnvPath())
langfuseFile := filepath.Join(envDir, LangfuseComposeFile)
c.LangfuseExtracted = checkFileExists(langfuseFile)
if h.dockerClient != nil {
workerExists, workerRunning := checkContainerExists(ctx, h.dockerClient, LangfuseWorkerContainerName)
webExists, webRunning := checkContainerExists(ctx, h.dockerClient, LangfuseWebContainerName)
c.LangfuseInstalled = workerExists && webExists
c.LangfuseRunning = workerRunning && webRunning
// check if langfuse-related volumes exist (indicates previous installation)
langfuseVolumes := []string{"langfuse-postgres-data", "langfuse-clickhouse-data", "langfuse-minio-data"}
c.LangfuseVolumesExist = checkVolumesExist(ctx, h.dockerClient, langfuseVolumes)
}
return nil
}
func (h *defaultCheckHandler) GatherObservabilityInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
otelHost := getEnvVar(h.appState, "OTEL_HOST", "")
c.ObservabilityConnected = otelHost != ""
c.ObservabilityExternal = otelHost != DefaultObservabilityEndpoint
envDir := filepath.Dir(h.appState.GetEnvPath())
obsFile := filepath.Join(envDir, ObservabilityComposeFile)
c.ObservabilityExtracted = checkFileExists(obsFile)
if h.dockerClient != nil {
exists, running := checkContainerExists(ctx, h.dockerClient, OpenTelemetryContainerName)
c.ObservabilityInstalled = exists
c.ObservabilityRunning = running
}
return nil
}
func (h *defaultCheckHandler) GatherSystemInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
// CPU check and count
c.SysCPUCount = runtime.NumCPU()
c.SysCPUOK = checkCPUResources()
// memory check and calculations
needsForPentagi, needsForGraphiti, needsForLangfuse, needsForObservability := determineComponentNeeds(c)
// calculate required memory using shared function
c.SysMemoryRequired = calculateRequiredMemoryGB(needsForPentagi, needsForGraphiti, needsForLangfuse, needsForObservability)
// get available memory and check if sufficient
c.SysMemoryAvailable = getAvailableMemoryGB()
c.SysMemoryOK = checkMemoryResources(needsForPentagi, needsForGraphiti, needsForLangfuse, needsForObservability)
// disk check and calculations
localComponents := countLocalComponentsToInstall(
c.PentagiInstalled,
c.GraphitiConnected, c.GraphitiExternal, c.GraphitiInstalled,
c.LangfuseConnected, c.LangfuseExternal, c.LangfuseInstalled,
c.ObservabilityConnected, c.ObservabilityExternal, c.ObservabilityInstalled,
)
// calculate required disk space using shared function
c.SysDiskRequired = calculateRequiredDiskGB(c.WorkerImageExists, localComponents)
// get available disk space and check if sufficient
c.SysDiskAvailable = getAvailableDiskGB(ctx)
c.SysDiskFreeSpaceOK = checkDiskSpaceWithContext(
ctx,
c.WorkerImageExists,
c.PentagiInstalled,
c.GraphitiConnected,
c.GraphitiExternal,
c.GraphitiInstalled,
c.LangfuseConnected,
c.LangfuseExternal,
c.LangfuseInstalled,
c.ObservabilityConnected,
c.ObservabilityExternal,
c.ObservabilityInstalled,
)
// network check with proxy and docker clients
proxyURL := getProxyURL(h.appState)
c.SysNetworkFailures = getNetworkFailures(ctx, proxyURL, h.dockerClient, h.workerClient)
c.SysNetworkOK = len(c.SysNetworkFailures) == 0
return nil
}
func (h *defaultCheckHandler) GatherUpdatesInfo(ctx context.Context, c *CheckResult) error {
h.mx.Lock()
defer h.mx.Unlock()
proxyURL := getProxyURL(h.appState)
updateServerURL := getEnvVar(h.appState, "UPDATE_SERVER_URL", DefaultUpdateServerEndpoint)
request := CheckUpdatesRequest{
InstallerOsType: runtime.GOOS,
InstallerVersion: InstallerVersion,
GraphitiConnected: c.GraphitiConnected,
GraphitiExternal: c.GraphitiExternal,
GraphitiInstalled: c.GraphitiInstalled,
LangfuseConnected: c.LangfuseConnected,
LangfuseExternal: c.LangfuseExternal,
LangfuseInstalled: c.LangfuseInstalled,
ObservabilityConnected: c.ObservabilityConnected,
ObservabilityExternal: c.ObservabilityExternal,
ObservabilityInstalled: c.ObservabilityInstalled,
}
// get PentAGI container image info
if h.dockerClient != nil && c.PentagiInstalled {
if imageInfo := getContainerImageInfo(ctx, h.dockerClient, PentagiContainerName); imageInfo != nil {
request.PentagiImageName = &imageInfo.Name
request.PentagiImageTag = &imageInfo.Tag
request.PentagiImageHash = &imageInfo.Hash
}
}
// get Worker image info from environment
if h.workerClient != nil {
defaultImage := getEnvVar(h.appState, "DOCKER_DEFAULT_IMAGE_FOR_PENTEST", DefaultImageForPentest)
if imageInfo := getImageInfo(ctx, h.workerClient, defaultImage); imageInfo != nil {
request.WorkerImageName = &imageInfo.Name
request.WorkerImageTag = &imageInfo.Tag
request.WorkerImageHash = &imageInfo.Hash
}
}
// get Graphiti image info if installed locally
if h.dockerClient != nil && c.GraphitiConnected && !c.GraphitiExternal && c.GraphitiInstalled {
if graphitiInfo := getContainerImageInfo(ctx, h.dockerClient, GraphitiContainerName); graphitiInfo != nil {
request.GraphitiImageName = &graphitiInfo.Name
request.GraphitiImageTag = &graphitiInfo.Tag
request.GraphitiImageHash = &graphitiInfo.Hash
}
if neo4jInfo := getContainerImageInfo(ctx, h.dockerClient, Neo4jContainerName); neo4jInfo != nil {
request.Neo4jImageName = &neo4jInfo.Name
request.Neo4jImageTag = &neo4jInfo.Tag
request.Neo4jImageHash = &neo4jInfo.Hash
}
}
// get Langfuse image info if installed locally
if h.dockerClient != nil && c.LangfuseConnected && !c.LangfuseExternal && c.LangfuseInstalled {
if workerInfo := getContainerImageInfo(ctx, h.dockerClient, LangfuseWorkerContainerName); workerInfo != nil {
request.LangfuseWorkerImageName = &workerInfo.Name
request.LangfuseWorkerImageTag = &workerInfo.Tag
request.LangfuseWorkerImageHash = &workerInfo.Hash
}
if webInfo := getContainerImageInfo(ctx, h.dockerClient, LangfuseWebContainerName); webInfo != nil {
request.LangfuseWebImageName = &webInfo.Name
request.LangfuseWebImageTag = &webInfo.Tag
request.LangfuseWebImageHash = &webInfo.Hash
}
}
// get Grafana and OpenTelemetry image info if observability installed locally
if h.dockerClient != nil && c.ObservabilityConnected && !c.ObservabilityExternal && c.ObservabilityInstalled {
if grafanaInfo := getContainerImageInfo(ctx, h.dockerClient, GrafanaContainerName); grafanaInfo != nil {
request.GrafanaImageName = &grafanaInfo.Name
request.GrafanaImageTag = &grafanaInfo.Tag
request.GrafanaImageHash = &grafanaInfo.Hash
}
if otelInfo := getContainerImageInfo(ctx, h.dockerClient, OpenTelemetryContainerName); otelInfo != nil {
request.OpenTelemetryImageName = &otelInfo.Name
request.OpenTelemetryImageTag = &otelInfo.Tag
request.OpenTelemetryImageHash = &otelInfo.Hash
}
}
response := checkUpdatesServer(ctx, updateServerURL, proxyURL, request)
if response != nil {
c.UpdateServerAccessible = true
c.InstallerIsUpToDate = response.InstallerIsUpToDate
c.PentagiIsUpToDate = response.PentagiIsUpToDate
c.GraphitiIsUpToDate = response.GraphitiIsUpToDate
c.LangfuseIsUpToDate = response.LangfuseIsUpToDate
c.ObservabilityIsUpToDate = response.ObservabilityIsUpToDate
c.WorkerIsUpToDate = response.WorkerIsUpToDate
} else {
c.UpdateServerAccessible = false
c.InstallerIsUpToDate = false
c.PentagiIsUpToDate = false
c.GraphitiIsUpToDate = false
c.LangfuseIsUpToDate = false
c.ObservabilityIsUpToDate = false
}
return nil
}
func Gather(ctx context.Context, appState state.State) (CheckResult, error) {
if appState == nil {
return CheckResult{}, ErrAppStateNotInitialized
}
c := CheckResult{
// default to the built-in handler
handler: &defaultCheckHandler{
mx: &sync.Mutex{},
appState: appState,
},
}
if err := c.GatherAllInfo(ctx); err != nil {
return c, err
}
return c, nil
}
func GatherWithHandler(ctx context.Context, handler CheckHandler) (CheckResult, error) {
if handler == nil {
return CheckResult{}, ErrHandlerNotInitialized
}
c := CheckResult{
handler: handler,
}
if err := handler.GatherAllInfo(ctx, &c); err != nil {
return c, err
}
return c, nil
}
backend/cmd/installer/checker/helpers.go
+1024
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/checker/helpers_test.go
+839
View File
@@ -0,0 +1,839 @@
package checker
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"time"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/state"
)
type mockState struct {
vars map[string]loader.EnvVar
envPath string
}
func (m *mockState) GetVar(key string) (loader.EnvVar, bool) {
if val, exists := m.vars[key]; exists {
return val, true
}
return loader.EnvVar{}, false
}
func (m *mockState) GetVars(names []string) (map[string]loader.EnvVar, map[string]bool) {
return m.vars, make(map[string]bool, len(names))
}
func (m *mockState) GetEnvPath() string {
return m.envPath
}
func (m *mockState) Exists() bool { return true }
func (m *mockState) Reset() error { return nil }
func (m *mockState) Commit() error { return nil }
func (m *mockState) IsDirty() bool { return false }
func (m *mockState) GetEulaConsent() bool { return true }
func (m *mockState) SetEulaConsent() error { return nil }
func (m *mockState) SetStack(stack []string) error { return nil }
func (m *mockState) GetStack() []string { return []string{} }
func (m *mockState) SetVar(name, value string) error { return nil }
func (m *mockState) ResetVar(name string) error { return nil }
func (m *mockState) SetVars(vars map[string]string) error { return nil }
func (m *mockState) ResetVars(names []string) error { return nil }
func (m *mockState) GetAllVars() map[string]loader.EnvVar { return m.vars }
func TestCheckFileExistsAndReadable(t *testing.T) {
f, err := os.CreateTemp("", "testfile")
if err != nil {
t.Fatal(err)
}
defer os.Remove(f.Name())
defer f.Close()
if !checkFileExists(f.Name()) {
t.Errorf("file should exist")
}
if !checkFileIsReadable(f.Name()) {
t.Errorf("file should be readable")
}
os.Remove(f.Name())
if checkFileExists(f.Name()) {
t.Errorf("file should not exist")
}
if checkFileIsReadable(f.Name()) {
t.Errorf("removed file should not be readable")
}
if checkFileExists("") {
t.Errorf("empty path should not exist")
}
if checkFileExists("/nonexistent/path/file.txt") {
t.Errorf("nonexistent file should not exist")
}
}
func TestGetEnvVar(t *testing.T) {
tests := []struct {
name string
vars map[string]loader.EnvVar
key string
defaultValue string
expected string
}{
{
name: "existing variable",
vars: map[string]loader.EnvVar{"FOO": {Value: "bar"}},
key: "FOO",
defaultValue: "default",
expected: "bar",
},
{
name: "non-existing variable",
vars: map[string]loader.EnvVar{},
key: "MISSING",
defaultValue: "default",
expected: "default",
},
{
name: "empty variable value",
vars: map[string]loader.EnvVar{"EMPTY": {Value: ""}},
key: "EMPTY",
defaultValue: "default",
expected: "default",
},
{
name: "nil state",
vars: nil,
key: "ANY",
defaultValue: "default",
expected: "default",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var appState state.State
if tt.vars != nil {
appState = &mockState{vars: tt.vars}
}
result := getEnvVar(appState, tt.key, tt.defaultValue)
if result != tt.expected {
t.Errorf("getEnvVar() = %q, want %q", result, tt.expected)
}
})
}
}
func TestExtractVersionFromOutput(t *testing.T) {
tests := []struct {
input string
expected string
}{
{"docker-compose version 1.29.2, build 5becea4c", "1.29.2"},
{"Docker Compose version v2.12.2", "2.12.2"},
{"Docker version 20.10.8, build 3967b7d", "20.10.8"},
{"no version here", ""},
{"v1.0.0-alpha", "1.0.0"},
{"version: 3.14.159", "3.14.159"},
{"", ""},
}
for _, tt := range tests {
t.Run(fmt.Sprintf("input_%s", tt.input), func(t *testing.T) {
result := extractVersionFromOutput(tt.input)
if result != tt.expected {
t.Errorf("extractVersionFromOutput(%q) = %q, want %q", tt.input, result, tt.expected)
}
})
}
}
func TestCheckDockerComposeVersionWithRunner(t *testing.T) {
t.Run("uses docker compose v2 output", func(t *testing.T) {
calls := 0
result := checkDockerComposeVersionWithRunner(func(name string, args ...string) ([]byte, error) {
calls++
if name != "docker" {
t.Fatalf("unexpected command %q", name)
}
if len(args) != 2 || args[0] != "compose" || args[1] != "version" {
t.Fatalf("unexpected args: %v", args)
}
return []byte("Docker Compose version v2.12.2"), nil
})
if calls != 1 {
t.Fatalf("expected 1 command invocation, got %d", calls)
}
if result.Version != "2.12.2" {
t.Fatalf("expected version 2.12.2, got %q", result.Version)
}
if !result.Valid {
t.Fatal("expected docker compose version to be valid")
}
})
t.Run("parses version from stdout even when error is returned", func(t *testing.T) {
calls := 0
result := checkDockerComposeVersionWithRunner(func(name string, args ...string) ([]byte, error) {
calls++
return []byte("Docker Compose version v2.12.2"), errors.New("exit status 1")
})
if calls != 1 {
t.Fatalf("expected 1 command invocation, got %d", calls)
}
if result.Version != "2.12.2" {
t.Fatalf("expected version 2.12.2, got %q", result.Version)
}
if !result.Valid {
t.Fatal("expected docker compose version to remain valid when stdout is parseable")
}
})
t.Run("fails when docker compose is unavailable", func(t *testing.T) {
calls := 0
result := checkDockerComposeVersionWithRunner(func(name string, args ...string) ([]byte, error) {
calls++
return nil, errors.New("executable file not found")
})
if calls != 1 {
t.Fatalf("expected 1 command invocation, got %d", calls)
}
if result.Version != "" {
t.Fatalf("expected empty version, got %q", result.Version)
}
if result.Valid {
t.Fatal("expected docker compose check to be invalid")
}
})
}
func TestCheckVersionCompatibility(t *testing.T) {
tests := []struct {
version string
minVersion string
expected bool
}{
{"1.2.3", "1.2.0", true},
{"1.2.0", "1.2.0", true},
{"1.1.9", "1.2.0", false},
{"2.0.0", "1.9.9", true},
{"1.2.3", "1.2.4", false},
{"", "1.0.0", false},
{"1.0.0", "", false},
{"invalid", "1.0.0", false},
{"1.0.0", "invalid", false},
{"1.2", "1.2.0", false}, // fewer parts should fail
{"1.2.0", "1.2", true}, // more parts should pass
}
for _, tt := range tests {
t.Run(fmt.Sprintf("%s_vs_%s", tt.version, tt.minVersion), func(t *testing.T) {
result := checkVersionCompatibility(tt.version, tt.minVersion)
if result != tt.expected {
t.Errorf("checkVersionCompatibility(%q, %q) = %v, want %v",
tt.version, tt.minVersion, result, tt.expected)
}
})
}
}
func TestParseImageRef(t *testing.T) {
tests := []struct {
imageRef string
imageID string
wantName string
wantTag string
wantHash string
}{
{"alpine:3.18", "sha256:abc", "alpine", "3.18", "sha256:abc"},
{"nginx", "", "nginx", "latest", ""},
{"nginx", "sha256:def", "nginx", "latest", "sha256:def"},
{"repo/nginx:1.2", "", "repo/nginx", "1.2", ""},
{"docker.io/library/ubuntu:latest", "", "library/ubuntu", "latest", ""},
{"nginx@sha256:deadbeef", "", "nginx", "latest", "sha256:deadbeef"},
{"myreg:5000/foo/bar:tag@sha256:beef", "", "foo/bar", "tag", "sha256:beef"},
{"localhost:5000/myapp:v1.0", "", "myapp", "v1.0", ""},
{"registry.example.com/team/app", "", "team/app", "latest", ""},
{"", "", "", "", ""},
{"ubuntu:", "", "ubuntu", "latest", ""},
{"ubuntu:@sha256:hash", "", "ubuntu", "latest", "sha256:hash"},
}
for _, tt := range tests {
t.Run(fmt.Sprintf("parse_%s", tt.imageRef), func(t *testing.T) {
if tt.imageRef == "" {
info := parseImageRef(tt.imageRef, tt.imageID)
if info != nil {
t.Errorf("parseImageRef(%q) should return nil for empty input", tt.imageRef)
}
return
}
info := parseImageRef(tt.imageRef, tt.imageID)
if info == nil {
t.Errorf("parseImageRef(%q) = nil, want non-nil", tt.imageRef)
return
}
// note: current implementation has some edge cases with registry parsing
// we test for non-nil result and basic structure rather than exact parsing
if info.Name == "" {
t.Errorf("parseImageRef(%q).Name should not be empty", tt.imageRef)
}
if info.Tag == "" {
t.Errorf("parseImageRef(%q).Tag should not be empty", tt.imageRef)
}
// hash may be empty, that's OK
})
}
}
func TestCheckCPUResources(t *testing.T) {
result := checkCPUResources()
// assuming test machine has at least 2 CPUs, this is reasonable for CI/dev environments
if !result {
t.Logf("CPU check returned false - this is expected on machines with < 2 CPUs")
}
}
func TestCheckMemoryResources(t *testing.T) {
tests := []struct {
name string
needsForPentagi bool
needsForGraphiti bool
needsForLangfuse bool
needsForObservability bool
expectMinimumRequirement bool
}{
{
name: "no components needed",
needsForPentagi: false,
needsForGraphiti: false,
needsForLangfuse: false,
needsForObservability: false,
expectMinimumRequirement: true,
},
{
name: "pentagi only",
needsForPentagi: true,
needsForGraphiti: false,
needsForLangfuse: false,
needsForObservability: false,
expectMinimumRequirement: false, // requires actual memory check
},
{
name: "all components",
needsForPentagi: true,
needsForGraphiti: true,
needsForLangfuse: true,
needsForObservability: true,
expectMinimumRequirement: false, // requires actual memory check
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := checkMemoryResources(tt.needsForPentagi, tt.needsForGraphiti, tt.needsForLangfuse, tt.needsForObservability)
if tt.expectMinimumRequirement && !result {
t.Errorf("checkMemoryResources() should return true when no components are needed")
}
// note: we can't reliably test memory checks across different environments
// the function will work correctly based on actual system memory
})
}
}
func TestCheckDiskSpaceWithContext(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
workerImageExists bool
pentagiInstalled bool
graphitiConnected bool
graphitiExternal bool
graphitiInstalled bool
langfuseConnected bool
langfuseExternal bool
langfuseInstalled bool
obsConnected bool
obsExternal bool
obsInstalled bool
expectHighSpace bool // whether we expect it to require more disk space
}{
{
name: "all installed and running",
workerImageExists: true,
pentagiInstalled: true,
graphitiConnected: true,
graphitiExternal: false,
graphitiInstalled: true,
langfuseConnected: true,
langfuseExternal: false,
langfuseInstalled: true,
obsConnected: true,
obsExternal: false,
obsInstalled: true,
expectHighSpace: false, // minimal space needed
},
{
name: "no worker images",
workerImageExists: false,
pentagiInstalled: true,
expectHighSpace: true, // needs to download images
},
{
name: "pentagi not installed",
workerImageExists: true,
pentagiInstalled: false,
expectHighSpace: false, // moderate space for components
},
{
name: "langfuse local not installed",
workerImageExists: true,
pentagiInstalled: true,
langfuseConnected: true,
langfuseExternal: false,
langfuseInstalled: false,
expectHighSpace: false, // moderate space for components
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := checkDiskSpaceWithContext(
ctx,
tt.workerImageExists,
tt.pentagiInstalled,
tt.graphitiConnected,
tt.graphitiExternal,
tt.graphitiInstalled,
tt.langfuseConnected,
tt.langfuseExternal,
tt.langfuseInstalled,
tt.obsConnected,
tt.obsExternal,
tt.obsInstalled,
)
// note: actual disk space check depends on OS and available space
// we mainly test that the function doesn't panic and returns a boolean
_ = result
})
}
}
func TestCheckUpdatesServer(t *testing.T) {
// test successful response
t.Run("successful_response", func(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
w.WriteHeader(http.StatusMethodNotAllowed)
return
}
if r.Header.Get("Content-Type") != "application/json" {
w.WriteHeader(http.StatusBadRequest)
return
}
if r.Header.Get("User-Agent") != UserAgent {
w.WriteHeader(http.StatusBadRequest)
return
}
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{
"installer_is_up_to_date": true,
"pentagi_is_up_to_date": false,
"langfuse_is_up_to_date": true,
"observability_is_up_to_date": false,
"worker_is_up_to_date": true
}`)
}))
defer ts.Close()
ctx := context.Background()
request := CheckUpdatesRequest{
InstallerVersion: "1.0.0",
InstallerOsType: "darwin",
}
response := checkUpdatesServer(ctx, ts.URL, "", request)
if response == nil {
t.Fatal("expected non-nil response")
}
if !response.InstallerIsUpToDate {
t.Error("expected installer to be up to date")
}
if response.PentagiIsUpToDate {
t.Error("expected pentagi to not be up to date")
}
if !response.LangfuseIsUpToDate {
t.Error("expected langfuse to be up to date")
}
if response.ObservabilityIsUpToDate {
t.Error("expected observability to not be up to date")
}
if !response.WorkerIsUpToDate {
t.Error("expected worker to be up to date")
}
})
// test server error
t.Run("server_error", func(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
}))
defer ts.Close()
ctx := context.Background()
request := CheckUpdatesRequest{InstallerVersion: "1.0.0"}
response := checkUpdatesServer(ctx, ts.URL, "", request)
if response != nil {
t.Error("expected nil response for server error")
}
})
// test invalid JSON response
t.Run("invalid_json", func(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `invalid json`)
}))
defer ts.Close()
ctx := context.Background()
request := CheckUpdatesRequest{InstallerVersion: "1.0.0"}
response := checkUpdatesServer(ctx, ts.URL, "", request)
if response != nil {
t.Error("expected nil response for invalid JSON")
}
})
// test context timeout
t.Run("context_timeout", func(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
time.Sleep(100 * time.Millisecond) // delay response
w.WriteHeader(http.StatusOK)
}))
defer ts.Close()
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
request := CheckUpdatesRequest{InstallerVersion: "1.0.0"}
response := checkUpdatesServer(ctx, ts.URL, "", request)
if response != nil {
t.Error("expected nil response for timeout")
}
})
// test proxy configuration
t.Run("with_proxy", func(t *testing.T) {
// create a proxy server that just forwards requests
proxyTs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"installer_is_up_to_date": true, "pentagi_is_up_to_date": true, "langfuse_is_up_to_date": true, "observability_is_up_to_date": true}`)
}))
defer proxyTs.Close()
ctx := context.Background()
request := CheckUpdatesRequest{InstallerVersion: "1.0.0"}
// note: testing with actual proxy setup is complex in unit tests
// this mainly tests that proxy URL doesn't cause the function to panic
response := checkUpdatesServer(ctx, proxyTs.URL, "http://invalid-proxy:8080", request)
// response might be nil due to proxy connection failure, which is expected
_ = response
})
// test malformed server URL
t.Run("malformed_url", func(t *testing.T) {
ctx := context.Background()
request := CheckUpdatesRequest{InstallerVersion: "1.0.0"}
response := checkUpdatesServer(ctx, "://invalid-url", "", request)
if response != nil {
t.Error("expected nil response for malformed URL")
}
})
}
func TestCreateTempFileForTesting(t *testing.T) {
// helper test to ensure temp file creation works for other tests
tmpDir := os.TempDir()
testFile := filepath.Join(tmpDir, "checker_test_file")
// create test file
err := os.WriteFile(testFile, []byte("test content"), 0644)
if err != nil {
t.Fatal(err)
}
defer os.Remove(testFile)
// verify it exists and is readable
if !checkFileExists(testFile) {
t.Error("test file should exist")
}
if !checkFileIsReadable(testFile) {
t.Error("test file should be readable")
}
// note: directory readability behavior is platform-dependent
// so we skip this assertion
}
func TestConstants(t *testing.T) {
// test that critical constants are defined
if InstallerVersion == "" {
t.Error("InstallerVersion should not be empty")
}
if UserAgent == "" {
t.Error("UserAgent should not be empty")
}
if !strings.Contains(UserAgent, InstallerVersion) {
t.Error("UserAgent should contain InstallerVersion")
}
if DefaultUpdateServerEndpoint == "" {
t.Error("DefaultUpdateServerEndpoint should not be empty")
}
if UpdatesCheckEndpoint == "" {
t.Error("UpdatesCheckEndpoint should not be empty")
}
// test memory and disk constants are reasonable
if MinFreeMemGB <= 0 {
t.Error("MinFreeMemGB should be positive")
}
if MinFreeMemGBForPentagi <= 0 {
t.Error("MinFreeMemGBForPentagi should be positive")
}
if MinFreeDiskGB <= 0 {
t.Error("MinFreeDiskGB should be positive")
}
if MinFreeDiskGBForWorkerImages <= MinFreeDiskGB {
t.Error("MinFreeDiskGBForWorkerImages should be larger than MinFreeDiskGB")
}
}
func TestCheckImageExistsEdgeCases(t *testing.T) {
ctx := context.Background()
// test with nil client
result := checkImageExists(ctx, nil, "nginx:latest")
if result {
t.Error("checkImageExists should return false for nil client")
}
// test with empty image name
// note: we can't test with real Docker client in unit tests
// but we can test that the function handles edge cases gracefully
}
func TestGetImageInfoEdgeCases(t *testing.T) {
ctx := context.Background()
// test with nil client
result := getImageInfo(ctx, nil, "nginx:latest")
if result != nil {
t.Error("getImageInfo should return nil for nil client")
}
// test with empty image name
// again, testing without real Docker client
}
func TestCheckUpdatesRequestStructure(t *testing.T) {
// test that CheckUpdatesRequest can be marshaled to JSON
request := CheckUpdatesRequest{
InstallerOsType: "darwin",
InstallerVersion: "1.0.0",
LangfuseConnected: true,
LangfuseExternal: false,
ObservabilityConnected: true,
ObservabilityExternal: false,
}
result := fmt.Sprintf("%+v", request)
if result == "" {
t.Error("CheckUpdatesRequest should be formattable")
}
// test with pointer fields
imageName := "test-image"
imageTag := "latest"
imageHash := "sha256:abc123"
request.PentagiImageName = &imageName
request.PentagiImageTag = &imageTag
request.PentagiImageHash = &imageHash
result = fmt.Sprintf("%+v", request)
if result == "" {
t.Error("CheckUpdatesRequest with pointers should be formattable")
}
}
func TestImageInfoStructure(t *testing.T) {
// test ImageInfo struct
info := &ImageInfo{
Name: "nginx",
Tag: "latest",
Hash: "sha256:abc123",
}
if info.Name != "nginx" {
t.Error("ImageInfo.Name should be set correctly")
}
if info.Tag != "latest" {
t.Error("ImageInfo.Tag should be set correctly")
}
if info.Hash != "sha256:abc123" {
t.Error("ImageInfo.Hash should be set correctly")
}
}
func TestCheckVolumesExist(t *testing.T) {
// note: this test uses a mock volume list since we can't rely on real Docker client in unit tests
// in real scenarios, checkVolumesExist is called with actual Docker API client
// test with nil client
t.Run("nil_client", func(t *testing.T) {
ctx := context.Background()
volumeNames := []string{"test-volume"}
result := checkVolumesExist(ctx, nil, volumeNames)
if result {
t.Error("checkVolumesExist should return false for nil client")
}
})
// test with empty volume list
t.Run("empty_volume_list", func(t *testing.T) {
ctx := context.Background()
// we can't create a real client in unit tests, so we pass nil
// the function should handle empty list gracefully
result := checkVolumesExist(ctx, nil, []string{})
if result {
t.Error("checkVolumesExist should return false for empty volume list")
}
})
// note: testing actual volume matching requires Docker integration tests
// the function logic handles:
// 1. Exact match: "pentagi-data" matches "pentagi-data"
// 2. Compose prefix match: "pentagi-data" matches "pentagi_pentagi-data"
// 3. Compose prefix match: "pentagi-postgres-data" matches "myproject_pentagi-postgres-data"
//
// This ensures compatibility with Docker Compose project prefixes
}
// mockDockerVolume simulates Docker API volume structure for testing
type mockDockerVolume struct {
Name string
}
func TestCheckVolumesExist_MatchingLogic(t *testing.T) {
// unit test for the matching logic without Docker client
// simulates what checkVolumesExist does internally
tests := []struct {
name string
existingVolumes []string
searchVolumes []string
expected bool
description string
}{
{
name: "exact match",
existingVolumes: []string{"pentagi-data", "other-volume"},
searchVolumes: []string{"pentagi-data"},
expected: true,
description: "should match exact volume name",
},
{
name: "compose prefix match",
existingVolumes: []string{"pentagi_pentagi-data", "pentagi_pentagi-ssl"},
searchVolumes: []string{"pentagi-data"},
expected: true,
description: "should match volume with compose project prefix",
},
{
name: "arbitrary prefix match",
existingVolumes: []string{"myproject_pentagi-postgres-data", "other_volume"},
searchVolumes: []string{"pentagi-postgres-data"},
expected: true,
description: "should match volume with any compose prefix",
},
{
name: "no match",
existingVolumes: []string{"other-volume", "another-volume"},
searchVolumes: []string{"pentagi-data"},
expected: false,
description: "should not match when volume doesn't exist",
},
{
name: "partial name should not match",
existingVolumes: []string{"pentagi-data-backup", "my-pentagi-data"},
searchVolumes: []string{"pentagi-data"},
expected: false,
description: "should not match partial names without underscore separator",
},
{
name: "match multiple search volumes",
existingVolumes: []string{"proj_pentagi-data", "langfuse-data"},
searchVolumes: []string{"pentagi-data", "langfuse-data", "missing-volume"},
expected: true,
description: "should return true if any search volume matches",
},
{
name: "empty existing volumes",
existingVolumes: []string{},
searchVolumes: []string{"pentagi-data"},
expected: false,
description: "should return false when no volumes exist",
},
{
name: "multiple compose prefixes",
existingVolumes: []string{"proj1_vol1", "proj2_vol2", "pentagi_pentagi-ssl"},
searchVolumes: []string{"pentagi-ssl"},
expected: true,
description: "should find volume among multiple compose projects",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// simulate the matching logic from checkVolumesExist
result := false
for _, volumeName := range tt.searchVolumes {
for _, existingVolume := range tt.existingVolumes {
if existingVolume == volumeName || strings.HasSuffix(existingVolume, "_"+volumeName) {
result = true
break
}
}
if result {
break
}
}
if result != tt.expected {
t.Errorf("%s: got %v, want %v", tt.description, result, tt.expected)
}
})
}
}
backend/cmd/installer/files/.gitignore
+3
View File
@@ -0,0 +1,3 @@
fs.go
fs_test.go
fs/
backend/cmd/installer/files/files.go
+401
View File
@@ -0,0 +1,401 @@
//go:generate go run generate.go
package files
import (
"fmt"
"io"
"io/fs"
"os"
"path/filepath"
"runtime"
"strings"
)
// FileStatus represents file integrity status
type FileStatus string
const (
FileStatusMissing FileStatus = "missing" // file does not exist
FileStatusModified FileStatus = "modified" // file exists but differs from embedded
FileStatusOK FileStatus = "ok" // file exists and matches embedded
)
// Files provides access to embedded and filesystem files
type Files interface {
// GetContent returns file content from embedded FS or filesystem fallback
GetContent(name string) ([]byte, error)
// Exists checks if file/directory exists in embedded FS
Exists(name string) bool
// ExistsInFS checks if file/directory exists in real filesystem
ExistsInFS(name string) bool
// Stat returns file info from embedded FS or filesystem fallback
Stat(name string) (fs.FileInfo, error)
// Copy copies file/directory from embedded FS to real filesystem
// dst is target directory, src name is preserved
Copy(src, dst string, rewrite bool) error
// Check returns file status comparing embedded vs filesystem
Check(name string, workingDir string) FileStatus
// List returns all embedded files with given prefix
List(prefix string) ([]string, error)
}
// EmbeddedProvider interface for generated embedded filesystem
type EmbeddedProvider interface {
GetContent(name string) ([]byte, error)
Exists(name string) bool
Stat(name string) (fs.FileInfo, error)
Copy(src, dst string, rewrite bool) error
List(prefix string) ([]string, error)
CheckHash(name, workingDir string) (bool, error)
ExpectedMode(name string) (fs.FileMode, bool)
}
// embeddedProvider holds reference to generated embedded provider
var embeddedProvider EmbeddedProvider = nil
// shouldCheckPermissions returns true if OS supports meaningful file permission bits
func shouldCheckPermissions() bool {
// Windows doesn't support Unix-style permission bits (rwxrwxrwx)
// It only has read-only attribute which is not comparable
return runtime.GOOS != "windows"
}
// files implements Files interface with fallback logic
type files struct {
linksDir string
}
func NewFiles() Files {
return &files{
linksDir: "links",
}
}
// GetContent returns file content from embedded FS or filesystem fallback
func (f *files) GetContent(name string) ([]byte, error) {
var embeddedErr error
if embeddedProvider != nil {
if content, err := embeddedProvider.GetContent(name); err == nil {
return content, nil
} else {
embeddedErr = err
}
}
// try filesystem fallback only if links directory exists
if f.ExistsInFS(name) {
return f.getContentFromFS(name)
}
// return informative error if both methods failed
if embeddedProvider == nil {
return nil, fmt.Errorf("embedded provider not initialized and file not found in filesystem: %s", name)
}
if embeddedErr != nil {
return nil, fmt.Errorf("file not found in embedded FS (%w) and not accessible in filesystem (links/%s)", embeddedErr, name)
}
return nil, fmt.Errorf("file not found: %s", name)
}
// Exists checks if file/directory exists in embedded FS
func (f *files) Exists(name string) bool {
if embeddedProvider != nil {
return embeddedProvider.Exists(name)
}
return false
}
// ExistsInFS checks if file/directory exists in real filesystem
func (f *files) ExistsInFS(name string) bool {
path := filepath.Join(f.linksDir, name)
_, err := os.Stat(path)
return err == nil
}
// Stat returns file info from embedded FS or filesystem fallback
func (f *files) Stat(name string) (fs.FileInfo, error) {
var embeddedErr error
if embeddedProvider != nil {
if info, err := embeddedProvider.Stat(name); err == nil {
return info, nil
} else {
embeddedErr = err
}
}
// try filesystem fallback only if file exists
if f.ExistsInFS(name) {
return f.statFromFS(name)
}
// return informative error if both methods failed
if embeddedProvider == nil {
return nil, fmt.Errorf("embedded provider not initialized and file not found in filesystem: %s", name)
}
if embeddedErr != nil {
return nil, fmt.Errorf("file not found in embedded FS (%w) and not accessible in filesystem (links/%s)", embeddedErr, name)
}
return nil, fmt.Errorf("file not found: %s", name)
}
// Copy copies file/directory from embedded FS to real filesystem
func (f *files) Copy(src, dst string, rewrite bool) error {
var embeddedErr error
if embeddedProvider != nil {
if err := embeddedProvider.Copy(src, dst, rewrite); err == nil {
return nil
} else {
embeddedErr = err
}
}
// try filesystem fallback only if source exists
if f.ExistsInFS(src) {
return f.copyFromFS(src, dst, rewrite)
}
// return informative error if both methods failed
if embeddedProvider == nil {
return fmt.Errorf("embedded provider not initialized and file not found in filesystem: %s", src)
}
if embeddedErr != nil {
return fmt.Errorf("cannot copy from embedded FS (%w) and not accessible in filesystem (links/%s)", embeddedErr, src)
}
return fmt.Errorf("file not found: %s", src)
}
// Check returns file status comparing embedded vs filesystem
func (f *files) Check(name string, workingDir string) FileStatus {
targetPath := filepath.Join(workingDir, name)
// check if file exists in filesystem
if _, err := os.Stat(targetPath); os.IsNotExist(err) {
return FileStatusMissing
}
// try hash-based comparison first (more efficient)
if embeddedProvider != nil {
if hashMatch, err := embeddedProvider.CheckHash(name, workingDir); err == nil {
if hashMatch {
// hash matches, also verify permission bits if available and meaningful on this OS
if shouldCheckPermissions() {
if expectedMode, ok := embeddedProvider.ExpectedMode(name); ok {
fsInfo, err := os.Stat(targetPath)
if err != nil {
return FileStatusMissing
}
if fsInfo.Mode().Perm() != expectedMode.Perm() {
return FileStatusModified
}
}
}
return FileStatusOK
}
// hash didn't match but no error, so it's definitely modified
return FileStatusModified
}
// if hash check failed (file not in metadata, etc.), fall back to content comparison
}
// fallback to content comparison
embeddedContent, err := f.GetContent(name)
if err != nil {
// if embedded doesn't exist, filesystem file is OK by default
return FileStatusOK
}
// read filesystem content
fsContent, err := os.ReadFile(targetPath)
if err != nil {
// cannot read filesystem file, consider it missing
return FileStatusMissing
}
// compare contents
if string(embeddedContent) == string(fsContent) {
// also compare permission bits when using filesystem fallback (only on Unix-like systems)
if shouldCheckPermissions() {
if infoExpected, err := f.statFromFS(name); err == nil {
if infoFS, err := os.Stat(targetPath); err == nil {
if infoFS.Mode().Perm() != infoExpected.Mode().Perm() {
return FileStatusModified
}
}
}
}
return FileStatusOK
}
return FileStatusModified
}
// List returns all embedded files with given prefix
func (f *files) List(prefix string) ([]string, error) {
if embeddedProvider != nil {
return embeddedProvider.List(prefix)
}
// fallback to filesystem listing
return f.listFromFS(prefix)
}
// getContentFromFS reads file content from real filesystem
func (f *files) getContentFromFS(name string) ([]byte, error) {
path := filepath.Join(f.linksDir, name)
return os.ReadFile(path)
}
// statFromFS gets file info from real filesystem
func (f *files) statFromFS(name string) (fs.FileInfo, error) {
path := filepath.Join(f.linksDir, name)
return os.Stat(path)
}
// copyFromFS copies file/directory from links directory to destination
func (f *files) copyFromFS(src, dst string, rewrite bool) error {
srcPath := filepath.Join(f.linksDir, src)
dstPath := filepath.Join(dst, src)
srcInfo, err := os.Stat(srcPath)
if err != nil {
return err
}
if srcInfo.IsDir() {
return f.copyDirFromFS(srcPath, dstPath, rewrite)
}
return f.copyFileFromFS(srcPath, dstPath, rewrite)
}
// copyFileFromFS copies single file
func (f *files) copyFileFromFS(src, dst string, rewrite bool) error {
if !rewrite {
if _, err := os.Stat(dst); err == nil {
return &os.PathError{Op: "copy", Path: dst, Err: os.ErrExist}
}
}
// Ensure destination directory exists
if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil {
return err
}
// read source mode to preserve permissions
srcInfo, err := os.Stat(src)
if err != nil {
return err
}
srcFile, err := os.Open(src)
if err != nil {
return err
}
defer srcFile.Close()
dstFile, err := os.Create(dst)
if err != nil {
return err
}
defer dstFile.Close()
if _, err = io.Copy(dstFile, srcFile); err != nil {
return err
}
// apply original permissions (best effort on all platforms)
// on Windows this may not preserve Unix-style bits, but will preserve read-only attribute
if err := os.Chmod(dst, srcInfo.Mode().Perm()); err != nil {
// on windows chmod may fail, but file is already copied
// don't fail the entire operation, just log or ignore
if runtime.GOOS != "windows" {
return err
}
}
return nil
}
// copyDirFromFS copies directory recursively
func (f *files) copyDirFromFS(src, dst string, rewrite bool) error {
if !rewrite {
if _, err := os.Stat(dst); err == nil {
return &os.PathError{Op: "copy", Path: dst, Err: os.ErrExist}
}
}
return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
if err != nil {
return err
}
relPath, err := filepath.Rel(src, path)
if err != nil {
return err
}
dstPath := filepath.Join(dst, relPath)
if info.IsDir() {
return os.MkdirAll(dstPath, info.Mode())
}
return f.copyFileFromFS(path, dstPath, rewrite)
})
}
// listFromFS lists files from filesystem with given prefix
func (f *files) listFromFS(prefix string) ([]string, error) {
var files []string
basePath := filepath.Join(f.linksDir, prefix)
// check if prefix path exists
if _, err := os.Stat(basePath); os.IsNotExist(err) {
return files, nil
}
// normalize prefix to forward slashes for consistent comparison
normalizedPrefix := filepath.ToSlash(prefix)
err := filepath.Walk(f.linksDir, func(path string, info os.FileInfo, err error) error {
if err != nil {
return err
}
// skip directories
if info.IsDir() {
return nil
}
// get relative path from links directory
relPath, err := filepath.Rel(f.linksDir, path)
if err != nil {
return err
}
// normalize to forward slashes for consistent comparison with embedded FS
normalizedRelPath := filepath.ToSlash(relPath)
// check if path starts with prefix
if normalizedPrefix == "" || strings.HasPrefix(normalizedRelPath, normalizedPrefix) {
files = append(files, normalizedRelPath)
}
return nil
})
return files, err
}
backend/cmd/installer/files/files_test.go
+609
View File
@@ -0,0 +1,609 @@
package files
import (
"os"
"path/filepath"
"runtime"
"strings"
"testing"
)
// newTestFiles creates a Files instance for testing with a custom links directory
func newTestFiles(linksDir string) Files {
return &files{
linksDir: linksDir,
}
}
func TestNewFiles(t *testing.T) {
f := NewFiles()
if f == nil {
t.Fatal("NewFiles() returned nil")
}
}
func TestGetContent_FromFS(t *testing.T) {
// Create temporary test directory structure
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
content, err := f.GetContent("test.txt")
if err != nil {
t.Fatalf("GetContent() error = %v", err)
}
expected := "test content"
if string(content) != expected {
t.Errorf("GetContent() = %q, want %q", string(content), expected)
}
}
func TestExistsInFS(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
if !f.ExistsInFS("test.txt") {
t.Error("ExistsInFS() = false, want true for existing file")
}
if f.ExistsInFS("nonexistent.txt") {
t.Error("ExistsInFS() = true, want false for non-existent file")
}
}
func TestStat_FromFS(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
info, err := f.Stat("test.txt")
if err != nil {
t.Fatalf("Stat() error = %v", err)
}
if info.IsDir() {
t.Error("Stat() IsDir() = true, want false for file")
}
if info.Size() == 0 {
t.Error("Stat() Size() = 0, want > 0")
}
}
func TestCopy_File(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
copyDstDir := t.TempDir()
defer os.RemoveAll(copyDstDir)
f := newTestFiles(testLinksDir)
err := f.Copy("test.txt", copyDstDir, false)
if err != nil {
t.Fatalf("Copy() error = %v", err)
}
// Verify file was copied
copiedPath := filepath.Join(copyDstDir, "test.txt")
content, err := os.ReadFile(copiedPath)
if err != nil {
t.Fatalf("Failed to read copied file: %v", err)
}
expected := "test content"
if string(content) != expected {
t.Errorf("Copied file content = %q, want %q", string(content), expected)
}
}
func TestCopy_PreservesExecutable_FromFS(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
if err := os.MkdirAll(testLinksDir, 0755); err != nil {
t.Fatalf("failed to create links dir: %v", err)
}
// create source file with specific permissions
src := filepath.Join(testLinksDir, "run.sh")
if err := os.WriteFile(src, []byte("#!/bin/sh\necho hi\n"), 0755); err != nil {
t.Fatalf("failed to create exec file: %v", err)
}
// get actual source mode (may differ on Windows)
srcInfo, err := os.Stat(src)
if err != nil {
t.Fatalf("failed to stat source: %v", err)
}
expectedMode := srcInfo.Mode().Perm()
copyDstDir := t.TempDir()
defer os.RemoveAll(copyDstDir)
f := newTestFiles(testLinksDir)
if err := f.Copy("run.sh", copyDstDir, false); err != nil {
t.Fatalf("Copy() error = %v", err)
}
// verify permissions preserved (whatever they actually are on this OS)
copied := filepath.Join(copyDstDir, "run.sh")
info, err := os.Stat(copied)
if err != nil {
t.Fatalf("failed to stat copied: %v", err)
}
if info.Mode().Perm() != expectedMode {
t.Errorf("copied mode = %o, want %o (source permissions not preserved)", info.Mode().Perm(), expectedMode)
}
}
func TestCheck_DetectsPermissionMismatch_FromFS(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
if err := os.MkdirAll(testLinksDir, 0755); err != nil {
t.Fatalf("failed to create links dir: %v", err)
}
// create source file
src := filepath.Join(testLinksDir, "tool.sh")
if err := os.WriteFile(src, []byte("#!/bin/sh\necho tool\n"), 0755); err != nil {
t.Fatalf("failed to create exec file: %v", err)
}
// get actual source mode
srcInfo, err := os.Stat(src)
if err != nil {
t.Fatalf("failed to stat source: %v", err)
}
f := newTestFiles(testLinksDir)
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
if err := f.Copy("tool.sh", workingDir, false); err != nil {
t.Fatalf("Copy() error = %v", err)
}
target := filepath.Join(workingDir, "tool.sh")
// try to change permissions
newMode := os.FileMode(0644)
if runtime.GOOS == "windows" {
// on Windows, we can only toggle read-only bit
newMode = 0444 // read-only
}
if err := os.Chmod(target, newMode); err != nil {
t.Fatalf("failed to chmod: %v", err)
}
// verify permissions actually changed
targetInfo, err := os.Stat(target)
if err != nil {
t.Fatalf("failed to stat target: %v", err)
}
if targetInfo.Mode().Perm() == srcInfo.Mode().Perm() {
// permissions didn't change on this OS, skip the rest
t.Skipf("cannot change file permissions on this OS (from %o to %o, got %o)",
srcInfo.Mode().Perm(), newMode, targetInfo.Mode().Perm())
}
status := f.Check("tool.sh", workingDir)
// on Windows, Check() doesn't compare permissions (by design)
// so even if permissions changed, status will be OK
if runtime.GOOS == "windows" {
if status != FileStatusOK {
t.Errorf("Check() on Windows = %v, want %v (permissions not checked on Windows)", status, FileStatusOK)
}
} else {
// on Unix, permissions should be checked
if status != FileStatusModified {
t.Errorf("Check() perms mismatch = %v, want %v", status, FileStatusModified)
}
}
}
func TestCopy_Directory(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksWithDirInDir(t, testLinksDir)
copyDstDir := t.TempDir()
defer os.RemoveAll(copyDstDir)
f := newTestFiles(testLinksDir)
err := f.Copy("testdir", copyDstDir, false)
if err != nil {
t.Fatalf("Copy() error = %v", err)
}
// Verify directory structure was copied
copiedFile := filepath.Join(copyDstDir, "testdir", "nested.txt")
content, err := os.ReadFile(copiedFile)
if err != nil {
t.Fatalf("Failed to read copied nested file: %v", err)
}
expected := "nested content"
if string(content) != expected {
t.Errorf("Copied nested file content = %q, want %q", string(content), expected)
}
}
func TestCopy_WithoutRewrite(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
copyDstDir := t.TempDir()
defer os.RemoveAll(copyDstDir)
f := newTestFiles(testLinksDir)
// Create existing file
existingPath := filepath.Join(copyDstDir, "test.txt")
err := os.WriteFile(existingPath, []byte("existing"), 0644)
if err != nil {
t.Fatalf("Failed to create existing file: %v", err)
}
// Try to copy without rewrite
err = f.Copy("test.txt", copyDstDir, false)
if err == nil {
t.Error("Copy() without rewrite should fail for existing file")
}
}
func TestCopy_WithRewrite(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
copyDstDir := t.TempDir()
defer os.RemoveAll(copyDstDir)
f := newTestFiles(testLinksDir)
// Create existing file
existingPath := filepath.Join(copyDstDir, "test.txt")
err := os.WriteFile(existingPath, []byte("existing"), 0644)
if err != nil {
t.Fatalf("Failed to create existing file: %v", err)
}
// Copy with rewrite
err = f.Copy("test.txt", copyDstDir, true)
if err != nil {
t.Fatalf("Copy() with rewrite error = %v", err)
}
// Verify file was overwritten
content, err := os.ReadFile(existingPath)
if err != nil {
t.Fatalf("Failed to read overwritten file: %v", err)
}
expected := "test content"
if string(content) != expected {
t.Errorf("Overwritten file content = %q, want %q", string(content), expected)
}
}
func TestExists_WithoutEmbedded(t *testing.T) {
f := NewFiles()
// Without embedded provider, Exists should return false
if f.Exists("any.txt") {
t.Error("Exists() = true, want false when no embedded provider")
}
}
func TestCopy_FromEmbedded(t *testing.T) {
f := NewFiles()
// This test only runs if embedded provider is available
if !f.Exists("docker-compose.yml") {
t.Skip("Skipping embedded test - no embedded provider")
}
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
err := f.Copy("docker-compose.yml", tmpDir, false)
if err != nil {
t.Fatalf("Copy() from embedded error = %v", err)
}
// Verify file was copied from embedded FS
copiedPath := filepath.Join(tmpDir, "docker-compose.yml")
content, err := os.ReadFile(copiedPath)
if err != nil {
t.Fatalf("Failed to read copied file: %v", err)
}
if len(content) == 0 {
t.Error("Copied file is empty")
}
// Verify content matches what we get from embedded FS
embeddedContent, err := f.GetContent("docker-compose.yml")
if err != nil {
t.Fatalf("Failed to get embedded content: %v", err)
}
if string(content) != string(embeddedContent) {
t.Error("Copied file content doesn't match embedded content")
}
}
func TestCheck_Missing(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
status := f.Check("test.txt", workingDir)
if status != FileStatusMissing {
t.Errorf("Check() = %v, want %v for missing file", status, FileStatusMissing)
}
}
func TestCheck_OK(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
// copy file to working directory
err := f.Copy("test.txt", workingDir, false)
if err != nil {
t.Fatalf("Copy() error = %v", err)
}
status := f.Check("test.txt", workingDir)
if status != FileStatusOK {
t.Errorf("Check() = %v, want %v for matching file", status, FileStatusOK)
}
}
func TestCheck_Modified(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
// create modified file in working directory
modifiedPath := filepath.Join(workingDir, "test.txt")
err := os.WriteFile(modifiedPath, []byte("modified content"), 0644)
if err != nil {
t.Fatalf("Failed to create modified file: %v", err)
}
status := f.Check("test.txt", workingDir)
if status != FileStatusModified {
t.Errorf("Check() = %v, want %v for modified file", status, FileStatusModified)
}
}
func TestList(t *testing.T) {
// test with real embedded provider (if available)
f := NewFiles()
// test listing with observability prefix (should exist in embedded)
files, err := f.List("observability")
if err != nil {
t.Fatalf("List() error = %v", err)
}
// we should have at least some observability files
if len(files) == 0 {
t.Error("List() with 'observability' prefix returned no files from embedded")
}
// verify we get some expected files
foundObservabilityFile := false
for _, file := range files {
if strings.HasPrefix(file, "observability/") {
foundObservabilityFile = true
break
}
}
if !foundObservabilityFile {
t.Error("List() with 'observability' prefix did not include any observability files")
}
// test listing with non-existent prefix
emptyFiles, err := f.List("nonexistent-prefix")
if err != nil {
t.Fatalf("List() with non-existent prefix error = %v", err)
}
if len(emptyFiles) != 0 {
t.Errorf("List() with non-existent prefix returned %d files, want 0", len(emptyFiles))
}
}
func TestList_NonExistentPrefix(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
testLinksDir := filepath.Join(tmpDir, "links")
setupTestLinksInDir(t, testLinksDir)
f := newTestFiles(testLinksDir)
files, err := f.List("nonexistent")
if err != nil {
t.Fatalf("List() error = %v", err)
}
if len(files) != 0 {
t.Errorf("List() for nonexistent prefix = %v, want empty slice", files)
}
}
func TestCheck_HashComparison_Embedded(t *testing.T) {
// test with real embedded files that have metadata
f := NewFiles()
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
// copy embedded file to working directory
embeddedFile := "docker-compose.yml"
err := f.Copy(embeddedFile, workingDir, false)
if err != nil {
t.Fatalf("Copy() error = %v", err)
}
// check should return OK (hash matches)
status := f.Check(embeddedFile, workingDir)
if status != FileStatusOK {
t.Errorf("Check() hash comparison = %v, want %v for embedded file", status, FileStatusOK)
}
}
func TestCheck_HashComparison_SameSize_DifferentContent(t *testing.T) {
// test case where file has same size but different content (different hash)
f := NewFiles()
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
// get metadata for docker-compose.yml to know its size
embeddedContent, err := f.GetContent("docker-compose.yml")
if err != nil {
t.Skip("Skipping test - docker-compose.yml not available")
}
originalSize := len(embeddedContent)
// create file with same size but different content
modifiedContent := make([]byte, originalSize)
for i := range modifiedContent {
modifiedContent[i] = 'X' // fill with different content
}
modifiedPath := filepath.Join(workingDir, "docker-compose.yml")
err = os.WriteFile(modifiedPath, modifiedContent, 0644)
if err != nil {
t.Fatalf("Failed to create modified file: %v", err)
}
// check should return Modified (same size, different hash)
status := f.Check("docker-compose.yml", workingDir)
if status != FileStatusModified {
t.Errorf("Check() same size different hash = %v, want %v", status, FileStatusModified)
}
}
func TestCheck_HashComparison_DifferentSize(t *testing.T) {
// test case where file has different size (quick size check should catch this)
f := NewFiles()
workingDir := t.TempDir()
defer os.RemoveAll(workingDir)
// create file with different size
modifiedContent := []byte("different size content")
modifiedPath := filepath.Join(workingDir, "docker-compose.yml")
err := os.WriteFile(modifiedPath, modifiedContent, 0644)
if err != nil {
t.Fatalf("Failed to create modified file: %v", err)
}
// check should return Modified (different size detected quickly)
status := f.Check("docker-compose.yml", workingDir)
if status != FileStatusModified {
t.Errorf("Check() different size = %v, want %v", status, FileStatusModified)
}
}
// Helper functions
// setupTestLinksInDir creates test files structure in specified directory
func setupTestLinksInDir(t *testing.T, linksDir string) {
err := os.MkdirAll(linksDir, 0755)
if err != nil {
t.Fatalf("Failed to create test links directory: %v", err)
}
testFile := filepath.Join(linksDir, "test.txt")
err = os.WriteFile(testFile, []byte("test content"), 0644)
if err != nil {
t.Fatalf("Failed to create test file: %v", err)
}
}
// setupTestLinksWithDirInDir creates test files and directories structure in specified directory
func setupTestLinksWithDirInDir(t *testing.T, linksDir string) {
setupTestLinksInDir(t, linksDir)
testDir := filepath.Join(linksDir, "testdir")
err := os.MkdirAll(testDir, 0755)
if err != nil {
t.Fatalf("Failed to create test directory: %v", err)
}
nestedFile := filepath.Join(testDir, "nested.txt")
err = os.WriteFile(nestedFile, []byte("nested content"), 0644)
if err != nil {
t.Fatalf("Failed to create nested test file: %v", err)
}
}
backend/cmd/installer/files/generate.go
+827
View File
@@ -0,0 +1,827 @@
//go:build ignore
// +build ignore
package main
import (
"crypto/sha256"
"encoding/json"
"fmt"
"io"
"log"
"os"
"path/filepath"
"runtime"
"strings"
)
// FileMetadata represents metadata for embedded files
type FileMetadata struct {
Path string `json:"path"`
Size int64 `json:"size"`
SHA256 string `json:"sha256"`
Mode uint32 `json:"mode"`
}
// MetadataFile contains all file metadata
type MetadataFile struct {
Files map[string]FileMetadata `json:"files"`
}
func main() {
linksDir := "links"
// check if links directory exists
if _, err := os.Stat(linksDir); os.IsNotExist(err) {
log.Printf("Links directory '%s' not found, skipping generation", linksDir)
return
}
// read all files and directories in links directory
entries, err := os.ReadDir(linksDir)
if err != nil {
log.Fatal(err)
}
var embedFiles []string
var fileContents = make(map[string]string)
metadata := MetadataFile{Files: make(map[string]FileMetadata)}
for _, entry := range entries {
entryPathRel := filepath.Join(linksDir, entry.Name())
entryPath, err := resolveSymlink(entryPathRel)
if err != nil {
log.Printf("Warning: could not resolve symlink %s: %v", entryPathRel, err)
continue
}
// follow symlinks to determine actual file type
info, err := os.Stat(entryPath)
if err != nil {
log.Printf("Warning: could not stat %s: %v", entryPath, err)
continue
}
if info.IsDir() {
// process directory recursively
// resolve symlink to get real directory path
realPath, err := evalSymlink(entryPath)
if err != nil {
log.Printf("Warning: could not resolve symlink %s: %v", entryPath, err)
continue
}
err = filepath.Walk(realPath, func(path string, walkInfo os.FileInfo, err error) error {
if err != nil {
return err
}
// skip directories themselves, only process files
if walkInfo.IsDir() {
return nil
}
// skip system files
if filepath.Base(path) == ".DS_Store" {
return nil
}
// get relative path from real directory root
relPathFromReal, err := filepath.Rel(realPath, path)
if err != nil {
return err
}
// construct relative path as it should appear in embedded fs
relPath := filepath.Join(entry.Name(), relPathFromReal)
embedFiles = append(embedFiles, relPath)
content, fileMeta, err := readFileContentWithMetadata(path)
if err != nil {
return err
}
relPath = strings.ReplaceAll(relPath, "\\", "/")
fileContents[relPath] = content
fileMeta.Path = relPath
metadata.Files[relPath] = fileMeta
return nil
})
if err != nil {
log.Fatal(err)
}
} else {
// process file
embedFiles = append(embedFiles, entry.Name())
content, fileMeta, err := readFileContentWithMetadata(entryPath)
if err != nil {
log.Printf("Warning: could not read file %s: %v", entryPath, err)
continue
}
fileContents[entry.Name()] = content
fileMeta.Path = entry.Name()
metadata.Files[entry.Name()] = fileMeta
}
}
// generate Go code for embedded provider
outputCode := `// Code generated by go generate; DO NOT EDIT.
package files
import (
"crypto/sha256"
"embed"
"encoding/json"
"fmt"
"io"
"io/fs"
"os"
"path/filepath"
"runtime"
"strings"
)
//go:embed fs/*
var embeddedFS embed.FS
// FileMetadata represents metadata for embedded files
type FileMetadata struct {
Path string ` + "`" + `json:"path"` + "`" + `
Size int64 ` + "`" + `json:"size"` + "`" + `
SHA256 string ` + "`" + `json:"sha256"` + "`" + `
Mode uint32 ` + "`" + `json:"mode"` + "`" + `
}
// MetadataFile contains all file metadata
type MetadataFile struct {
Files map[string]FileMetadata ` + "`" + `json:"files"` + "`" + `
}
// embeddedProvider implements EmbeddedProvider interface
type embeddedProviderImpl struct {
metadata *MetadataFile
}
func init() {
ep := &embeddedProviderImpl{}
// load metadata
if metaContent, err := embeddedFS.ReadFile("fs/.meta.json"); err == nil {
var meta MetadataFile
if err := json.Unmarshal(metaContent, &meta); err == nil {
ep.metadata = &meta
}
}
embeddedProvider = ep
}
// toEmbedPath converts OS-specific path to embed.FS compatible path (forward slashes)
func toEmbedPath(parts ...string) string {
return filepath.ToSlash(filepath.Join(parts...))
}
// GetContent returns file content from embedded filesystem
func (ep *embeddedProviderImpl) GetContent(name string) ([]byte, error) {
return embeddedFS.ReadFile(toEmbedPath("fs", name))
}
// Exists checks if file/directory exists in embedded filesystem
func (ep *embeddedProviderImpl) Exists(name string) bool {
_, err := fs.Stat(embeddedFS, toEmbedPath("fs", name))
return err == nil
}
// Stat returns file info from embedded filesystem
func (ep *embeddedProviderImpl) Stat(name string) (fs.FileInfo, error) {
return fs.Stat(embeddedFS, toEmbedPath("fs", name))
}
// Copy copies file/directory from embedded FS to real filesystem
func (ep *embeddedProviderImpl) Copy(src, dst string, rewrite bool) error {
srcPath := toEmbedPath("fs", src)
dstPath := filepath.Join(dst, src)
info, err := fs.Stat(embeddedFS, srcPath)
if err != nil {
return err
}
if info.IsDir() {
return ep.copyDirFromEmbed(srcPath, dstPath, rewrite)
}
return ep.copyFileFromEmbed(srcPath, dstPath, rewrite)
}
// copyFileFromEmbed copies single file from embedded FS using streaming
func (ep *embeddedProviderImpl) copyFileFromEmbed(src, dst string, rewrite bool) error {
info, err := os.Stat(dst)
if !rewrite && info != nil && err == nil {
return &os.PathError{Op: "copy", Path: dst, Err: os.ErrExist}
}
// if rewrite is true and destination is a directory, remove it to avoid errors
if rewrite && info != nil && info.IsDir() {
if err := os.RemoveAll(dst); err != nil {
return err
}
}
// ensure destination directory exists
if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil {
return err
}
// open embedded file for streaming
srcFile, err := embeddedFS.Open(src)
if err != nil {
return err
}
defer srcFile.Close()
// create destination file
dstFile, err := os.Create(dst)
if err != nil {
return err
}
defer dstFile.Close()
// stream copy without loading full file into memory
_, err = io.Copy(dstFile, srcFile)
if err != nil {
return err
}
// apply permissions if metadata available
if ep.metadata != nil {
// src has prefix "fs/"; strip to get metadata key
// normalize path separators for metadata lookup
rel := strings.TrimPrefix(filepath.ToSlash(src), "fs/")
if meta, ok := ep.metadata.Files[rel]; ok {
// best effort: try to apply permissions
// on Windows this may not work as expected for Unix-style permissions
// but will preserve read-only attribute
if chmodErr := os.Chmod(dst, fs.FileMode(meta.Mode)); chmodErr != nil {
// on Windows chmod may fail for some modes, but file is already copied
// don't fail the entire operation
if runtime.GOOS != "windows" {
return chmodErr
}
}
}
}
return nil
}
// copyDirFromEmbed copies directory recursively from embedded FS
func (ep *embeddedProviderImpl) copyDirFromEmbed(src, dst string, rewrite bool) error {
if !rewrite {
if _, err := os.Stat(dst); err == nil {
return &os.PathError{Op: "copy", Path: dst, Err: os.ErrExist}
}
}
return fs.WalkDir(embeddedFS, src, func(walkPath string, d fs.DirEntry, err error) error {
if err != nil {
return err
}
// embedded FS always uses forward slashes, even on Windows
// calculate relative path: walkPath is guaranteed to start with src
var relPath string
if walkPath == src {
// walking the root directory itself
relPath = ""
} else {
// walkPath = "fs/dir/file.txt", src = "fs/dir" → relPath = "file.txt"
relPath = strings.TrimPrefix(walkPath, src+"/")
}
// convert forward-slash path to OS-specific path for destination
dstPath := filepath.Join(dst, filepath.FromSlash(relPath))
if d.IsDir() {
return os.MkdirAll(dstPath, 0755)
}
return ep.copyFileFromEmbed(walkPath, dstPath, rewrite)
})
}
// List returns all embedded files with given prefix
func (ep *embeddedProviderImpl) List(prefix string) ([]string, error) {
var files []string
// normalize prefix to forward slashes for comparison with embedded FS paths
normalizedPrefix := filepath.ToSlash(prefix)
err := fs.WalkDir(embeddedFS, "fs", func(walkPath string, d fs.DirEntry, err error) error {
if err != nil {
return err
}
// skip directories
if d.IsDir() {
return nil
}
// embedded FS always uses forward slashes, even on Windows
// walkPath is guaranteed to start with "fs/" since we're walking from "fs"
// walkPath = "fs/dir/file.txt" → relPath = "dir/file.txt"
relPath := strings.TrimPrefix(walkPath, "fs/")
// check if path starts with prefix
if normalizedPrefix == "" || strings.HasPrefix(relPath, normalizedPrefix) {
files = append(files, relPath)
}
return nil
})
return files, err
}
// CheckHash compares file hash with embedded metadata
func (ep *embeddedProviderImpl) CheckHash(name, workingDir string) (bool, error) {
if ep.metadata == nil {
return false, fmt.Errorf("no metadata available")
}
// normalize path separators for metadata lookup
normalizedName := filepath.ToSlash(name)
meta, exists := ep.metadata.Files[normalizedName]
if !exists {
return false, fmt.Errorf("file not found in metadata")
}
targetPath := filepath.Join(workingDir, name)
// check file size first (quick check)
fsInfo, err := os.Stat(targetPath)
if err != nil {
return false, err
}
if fsInfo.Size() != meta.Size {
return false, nil // different size, definitely different
}
// calculate hash of filesystem file
fsFile, err := os.Open(targetPath)
if err != nil {
return false, err
}
defer fsFile.Close()
hash := sha256.New()
if _, err := io.Copy(hash, fsFile); err != nil {
return false, err
}
fsHash := fmt.Sprintf("%x", hash.Sum(nil))
return fsHash == meta.SHA256, nil
}
// ExpectedMode returns expected permission bits for a file from metadata
func (ep *embeddedProviderImpl) ExpectedMode(name string) (fs.FileMode, bool) {
if ep.metadata == nil {
return 0, false
}
// normalize path separators for metadata lookup
normalizedName := filepath.ToSlash(name)
meta, ok := ep.metadata.Files[normalizedName]
if !ok {
return 0, false
}
return fs.FileMode(meta.Mode), true
}
`
outputTests := `// Code generated by go generate; DO NOT EDIT.
package files
import (
"os"
"path/filepath"
"strings"
"testing"
)
// TestToEmbedPath verifies cross-platform path normalization
func TestToEmbedPath(t *testing.T) {
tests := []struct {
name string
parts []string
expect string
}{
{"simple", []string{"fs", ".env"}, "fs/.env"},
{"nested", []string{"fs", "observability", "grafana", "config.yml"}, "fs/observability/grafana/config.yml"},
{"single", []string{"docker-compose.yml"}, "docker-compose.yml"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := toEmbedPath(tt.parts...)
if result != tt.expect {
t.Errorf("toEmbedPath(%v) = %q, want %q", tt.parts, result, tt.expect)
}
// verify no backslashes (Windows compatibility)
if strings.Contains(result, "\\") {
t.Errorf("toEmbedPath() returned path with backslash: %q", result)
}
})
}
}
// TestEmbeddedProvider_PathNormalization tests that embedded provider works with both "/" and "\" in input
func TestEmbeddedProvider_PathNormalization(t *testing.T) {
if embeddedProvider == nil {
t.Skip("embedded provider not available")
}
// test with known embedded file
testCases := []struct {
name string
path string
wantFile string // expected file in embedded FS
}{
{"unix_style", ".env", ".env"},
{"unix_nested", "observability/grafana/config/grafana.ini", "observability/grafana/config/grafana.ini"},
{"windows_style", filepath.Join("observability", "grafana", "config", "grafana.ini"), "observability/grafana/config/grafana.ini"},
{"mixed_depth", filepath.Join("providers-configs", "deepseek.provider.yml"), "providers-configs/deepseek.provider.yml"},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// test Exists
if !embeddedProvider.Exists(tc.path) {
t.Errorf("Exists(%q) = false, want true", tc.path)
}
// test GetContent
content, err := embeddedProvider.GetContent(tc.path)
if err != nil {
t.Errorf("GetContent(%q) error = %v", tc.path, err)
}
if len(content) == 0 {
t.Errorf("GetContent(%q) returned empty content", tc.path)
}
// test Stat
info, err := embeddedProvider.Stat(tc.path)
if err != nil {
t.Errorf("Stat(%q) error = %v", tc.path, err)
}
if info.Size() != int64(len(content)) {
t.Errorf("Stat(%q).Size() = %d, want %d", tc.path, info.Size(), len(content))
}
})
}
}
// TestEmbeddedProvider_CheckHash tests hash verification with normalized paths
func TestEmbeddedProvider_CheckHash(t *testing.T) {
if embeddedProvider == nil {
t.Skip("embedded provider not available")
}
workingDir := t.TempDir()
testFiles := []string{
".env",
filepath.Join("observability", "loki", "config.yml"),
filepath.Join("providers-configs", "deepseek.provider.yml"),
}
for _, testFile := range testFiles {
t.Run(testFile, func(t *testing.T) {
// copy file to working directory
err := embeddedProvider.Copy(testFile, workingDir, false)
if err != nil {
t.Fatalf("Copy(%q) error = %v", testFile, err)
}
// verify hash matches
match, err := embeddedProvider.CheckHash(testFile, workingDir)
if err != nil {
t.Errorf("CheckHash(%q) error = %v", testFile, err)
}
if !match {
t.Errorf("CheckHash(%q) = false, want true for just copied file", testFile)
}
// verify ExpectedMode works
if mode, ok := embeddedProvider.ExpectedMode(testFile); !ok {
t.Errorf("ExpectedMode(%q) not found in metadata", testFile)
} else if mode == 0 {
t.Errorf("ExpectedMode(%q) = 0, want non-zero", testFile)
}
})
}
}
// TestEmbeddedProvider_Copy tests directory and file copying with path normalization
func TestEmbeddedProvider_Copy(t *testing.T) {
if embeddedProvider == nil {
t.Skip("embedded provider not available")
}
workingDir := t.TempDir()
testCases := []struct {
name string
src string
expectMin int // minimum files expected
}{
{"single_file", ".env", 1},
{"nested_file", filepath.Join("observability", "loki", "config.yml"), 1},
{"directory", filepath.Join("observability", "loki"), 1},
{"deep_directory", filepath.Join("observability", "grafana", "dashboards"), 3},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
dstDir := filepath.Join(workingDir, tc.name)
err := embeddedProvider.Copy(tc.src, dstDir, false)
if err != nil {
t.Fatalf("Copy(%q) error = %v", tc.src, err)
}
// verify file/directory exists at destination
dstPath := filepath.Join(dstDir, tc.src)
info, err := os.Stat(dstPath)
if err != nil {
t.Fatalf("Stat(%q) after copy error = %v", dstPath, err)
}
// for files, verify content matches
if !info.IsDir() {
embeddedContent, _ := embeddedProvider.GetContent(tc.src)
copiedContent, _ := os.ReadFile(dstPath)
if string(embeddedContent) != string(copiedContent) {
t.Errorf("copied file content differs from embedded")
}
} else {
// for directories, count files
var fileCount int
filepath.Walk(dstPath, func(path string, info os.FileInfo, err error) error {
if err == nil && !info.IsDir() {
fileCount++
}
return nil
})
if fileCount < tc.expectMin {
t.Errorf("copied directory has %d files, want at least %d", fileCount, tc.expectMin)
}
}
})
}
}
// TestEmbeddedProvider_List tests listing with various prefix formats
func TestEmbeddedProvider_List(t *testing.T) {
if embeddedProvider == nil {
t.Skip("embedded provider not available")
}
testCases := []struct {
name string
prefix string
expectMin int // minimum files expected
mustHave []string // paths that must be in results (normalized)
mustNotHave []string // paths that must not be in results
}{
{
name: "unix_style_prefix",
prefix: "observability/loki",
expectMin: 1,
mustHave: []string{"observability/loki/config.yml"},
},
{
name: "windows_style_prefix",
prefix: filepath.Join("observability", "grafana"),
expectMin: 5,
mustHave: []string{"observability/grafana/config/grafana.ini"},
},
{
name: "providers_prefix",
prefix: "providers-configs",
expectMin: 5,
mustHave: []string{"providers-configs/deepseek.provider.yml"},
mustNotHave: []string{"observability/loki/config.yml"},
},
{
name: "empty_prefix",
prefix: "",
expectMin: 20, // should return all files
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
files, err := embeddedProvider.List(tc.prefix)
if err != nil {
t.Fatalf("List(%q) error = %v", tc.prefix, err)
}
if len(files) < tc.expectMin {
t.Errorf("List(%q) returned %d files, want at least %d", tc.prefix, len(files), tc.expectMin)
}
// verify all returned paths use forward slashes
for _, f := range files {
if strings.Contains(f, "\\") {
t.Errorf("List(%q) returned path with backslash: %q", tc.prefix, f)
}
}
// verify must-have files are present
fileSet := make(map[string]bool)
for _, f := range files {
fileSet[f] = true
}
for _, mustHave := range tc.mustHave {
if !fileSet[mustHave] {
t.Errorf("List(%q) missing expected file: %q", tc.prefix, mustHave)
}
}
// verify must-not-have files are absent
for _, mustNotHave := range tc.mustNotHave {
if fileSet[mustNotHave] {
t.Errorf("List(%q) contains unexpected file: %q", tc.prefix, mustNotHave)
}
}
})
}
}
// TestEmbeddedProvider_PermissionsPreserved tests that file permissions are preserved on copy
func TestEmbeddedProvider_PermissionsPreserved(t *testing.T) {
if embeddedProvider == nil {
t.Skip("embedded provider not available")
}
workingDir := t.TempDir()
// find an executable file in embedded FS
executableFiles := []string{
"observability/jaeger/bin/jaeger-clickhouse-linux-amd64",
"observability/jaeger/bin/jaeger-clickhouse-linux-arm64",
}
for _, testFile := range executableFiles {
if !embeddedProvider.Exists(testFile) {
continue
}
t.Run(testFile, func(t *testing.T) {
expectedMode, ok := embeddedProvider.ExpectedMode(testFile)
if !ok {
t.Skip("no mode metadata available")
}
err := embeddedProvider.Copy(testFile, workingDir, false)
if err != nil {
t.Fatalf("Copy(%q) error = %v", testFile, err)
}
copiedPath := filepath.Join(workingDir, testFile)
info, err := os.Stat(copiedPath)
if err != nil {
t.Fatalf("Stat(%q) error = %v", copiedPath, err)
}
if info.Mode().Perm() != expectedMode.Perm() {
t.Errorf("copied file mode = %o, want %o", info.Mode().Perm(), expectedMode.Perm())
}
})
break // test only one executable file
}
}
`
// create fs directory
err = os.MkdirAll("fs", 0755)
if err != nil {
log.Fatal(err)
}
// copy file contents to fs directory
for filename, content := range fileContents {
destPath := filepath.Join("fs", filename)
// create directories if needed
err = os.MkdirAll(filepath.Dir(destPath), 0755)
if err != nil {
log.Fatal(err)
}
err = os.WriteFile(destPath, []byte(content), 0644)
if err != nil {
log.Fatal(err)
}
// apply source file permissions if present
if meta, ok := metadata.Files[filename]; ok {
if chmodErr := os.Chmod(destPath, os.FileMode(meta.Mode)); chmodErr != nil {
log.Fatal(chmodErr)
}
}
}
// write metadata file
metadataContent, err := json.MarshalIndent(metadata, "", " ")
if err != nil {
log.Fatal(err)
}
metaPath := filepath.Join("fs", ".meta.json")
err = os.WriteFile(metaPath, metadataContent, 0644)
if err != nil {
log.Fatal(err)
}
// write generated Go code
err = os.WriteFile("fs.go", []byte(outputCode), 0644)
if err != nil {
log.Fatal(err)
}
err = os.WriteFile("fs_test.go", []byte(outputTests), 0644)
if err != nil {
log.Fatal(err)
}
fmt.Printf("Generated embedded files for: %v\n", embedFiles)
}
func readSymlinkWindows(symlinkPath string) (string, error) {
fileContent, err := os.ReadFile(symlinkPath)
if err != nil {
return "", fmt.Errorf("failed to read symlink on windows")
}
content := strings.Split(string(fileContent), "\n")[0]
if contentLen := len(content); contentLen > 255 || contentLen == 0 {
return "", fmt.Errorf("invalid symlink path")
}
content = strings.ReplaceAll(content, "\\", "/")
content = filepath.Join(filepath.Dir(symlinkPath), content)
return filepath.Abs(content)
}
func resolveSymlink(entryPath string) (string, error) {
if runtime.GOOS == "windows" {
return readSymlinkWindows(entryPath)
}
return entryPath, nil
}
func evalSymlink(entryPath string) (string, error) {
if runtime.GOOS == "windows" {
return filepath.Abs(entryPath)
}
return filepath.EvalSymlinks(entryPath)
}
func readFileContentWithMetadata(filename string) (string, FileMetadata, error) {
file, err := os.Open(filename)
if err != nil {
return "", FileMetadata{}, err
}
defer file.Close()
// get file info for size
info, err := file.Stat()
if err != nil {
return "", FileMetadata{}, err
}
// calculate hash while reading content
hash := sha256.New()
teeReader := io.TeeReader(file, hash)
content, err := io.ReadAll(teeReader)
if err != nil {
return "", FileMetadata{}, err
}
meta := FileMetadata{
Size: info.Size(),
SHA256: fmt.Sprintf("%x", hash.Sum(nil)),
Mode: uint32(info.Mode().Perm()),
}
return string(content), meta, nil
}
backend/cmd/installer/files/links/.env
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../.env.example
backend/cmd/installer/files/links/EULA.md
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../EULA.md
backend/cmd/installer/files/links/docker-compose-graphiti.yml
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../docker-compose-graphiti.yml
backend/cmd/installer/files/links/docker-compose-langfuse.yml
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../docker-compose-langfuse.yml
backend/cmd/installer/files/links/docker-compose-observability.yml
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../docker-compose-observability.yml
backend/cmd/installer/files/links/docker-compose.yml
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../docker-compose.yml
backend/cmd/installer/files/links/example.custom.provider.yml
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../examples/configs/custom-openai.provider.yml
backend/cmd/installer/files/links/example.ollama.provider.yml
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../examples/configs/ollama-llama318b-instruct.provider.yml
backend/cmd/installer/files/links/observability
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../observability
backend/cmd/installer/files/links/providers-configs
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../../examples/configs
backend/cmd/installer/hardening/hardening.go
+366
View File
@@ -0,0 +1,366 @@
package hardening
import (
"crypto/rand"
"encoding/hex"
"fmt"
"net/url"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/state"
"github.com/google/uuid"
"github.com/vxcontrol/cloud/sdk"
"github.com/vxcontrol/cloud/system"
)
type HardeningArea string
const (
HardeningAreaPentagi HardeningArea = "pentagi"
HardeningAreaLangfuse HardeningArea = "langfuse"
HardeningAreaGraphiti HardeningArea = "graphiti"
)
type HardeningPolicyType string
const (
HardeningPolicyTypeDefault HardeningPolicyType = "default"
HardeningPolicyTypeHex HardeningPolicyType = "hex"
HardeningPolicyTypeUUID HardeningPolicyType = "uuid"
HardeningPolicyTypeBoolTrue HardeningPolicyType = "bool_true"
HardeningPolicyTypeBoolFalse HardeningPolicyType = "bool_false"
)
type HardeningPolicy struct {
Type HardeningPolicyType
Length int // length of the random string
Prefix string // prefix for the random string
}
var varsForHardening = map[HardeningArea][]string{
HardeningAreaPentagi: {
"COOKIE_SIGNING_SALT",
"PENTAGI_POSTGRES_PASSWORD",
"LOCAL_SCRAPER_USERNAME",
"LOCAL_SCRAPER_PASSWORD",
"SCRAPER_PRIVATE_URL",
},
HardeningAreaGraphiti: {
"NEO4J_PASSWORD",
},
HardeningAreaLangfuse: {
"LANGFUSE_POSTGRES_PASSWORD",
"LANGFUSE_CLICKHOUSE_PASSWORD",
"LANGFUSE_S3_ACCESS_KEY_ID",
"LANGFUSE_S3_SECRET_ACCESS_KEY",
"LANGFUSE_REDIS_AUTH",
"LANGFUSE_SALT",
"LANGFUSE_ENCRYPTION_KEY",
"LANGFUSE_NEXTAUTH_SECRET",
"LANGFUSE_INIT_PROJECT_ID",
"LANGFUSE_INIT_PROJECT_PUBLIC_KEY",
"LANGFUSE_INIT_PROJECT_SECRET_KEY",
"LANGFUSE_AUTH_DISABLE_SIGNUP",
"LANGFUSE_PROJECT_ID",
"LANGFUSE_PUBLIC_KEY",
"LANGFUSE_SECRET_KEY",
},
}
var varsForHardeningDefault = map[string]string{
"COOKIE_SIGNING_SALT": "salt",
"PENTAGI_POSTGRES_PASSWORD": "postgres",
"NEO4J_PASSWORD": "devpassword",
"LOCAL_SCRAPER_USERNAME": "someuser",
"LOCAL_SCRAPER_PASSWORD": "somepass",
"SCRAPER_PRIVATE_URL": "https://someuser:somepass@scraper/",
"LANGFUSE_POSTGRES_PASSWORD": "postgres",
"LANGFUSE_CLICKHOUSE_PASSWORD": "clickhouse",
"LANGFUSE_S3_ACCESS_KEY_ID": "accesskey",
"LANGFUSE_S3_SECRET_ACCESS_KEY": "secretkey",
"LANGFUSE_REDIS_AUTH": "redispassword",
"LANGFUSE_SALT": "salt",
"LANGFUSE_ENCRYPTION_KEY": "0000000000000000000000000000000000000000000000000000000000000000",
"LANGFUSE_NEXTAUTH_SECRET": "secret",
"LANGFUSE_INIT_PROJECT_ID": "cm47619l0000872mcd2dlbqwb",
"LANGFUSE_INIT_PROJECT_PUBLIC_KEY": "pk-lf-00000000-0000-0000-0000-000000000000",
"LANGFUSE_INIT_PROJECT_SECRET_KEY": "sk-lf-00000000-0000-0000-0000-000000000000",
"LANGFUSE_AUTH_DISABLE_SIGNUP": "false",
"LANGFUSE_PROJECT_ID": "",
"LANGFUSE_PUBLIC_KEY": "",
"LANGFUSE_SECRET_KEY": "",
}
var varsHardeningSyncLangfuse = map[string]string{
"LANGFUSE_PROJECT_ID": "LANGFUSE_INIT_PROJECT_ID",
"LANGFUSE_PUBLIC_KEY": "LANGFUSE_INIT_PROJECT_PUBLIC_KEY",
"LANGFUSE_SECRET_KEY": "LANGFUSE_INIT_PROJECT_SECRET_KEY",
}
var varsHardeningPolicies = map[HardeningArea]map[string]HardeningPolicy{
HardeningAreaPentagi: {
"COOKIE_SIGNING_SALT": {Type: HardeningPolicyTypeHex, Length: 32},
"PENTAGI_POSTGRES_PASSWORD": {Type: HardeningPolicyTypeDefault, Length: 18},
"LOCAL_SCRAPER_USERNAME": {Type: HardeningPolicyTypeDefault, Length: 10},
"LOCAL_SCRAPER_PASSWORD": {Type: HardeningPolicyTypeDefault, Length: 12},
// SCRAPER_PRIVATE_URL is handled specially in DoHardening logic
},
HardeningAreaGraphiti: {
"NEO4J_PASSWORD": {Type: HardeningPolicyTypeDefault, Length: 18},
},
HardeningAreaLangfuse: {
"LANGFUSE_POSTGRES_PASSWORD": {Type: HardeningPolicyTypeDefault, Length: 18},
"LANGFUSE_CLICKHOUSE_PASSWORD": {Type: HardeningPolicyTypeDefault, Length: 18},
"LANGFUSE_S3_ACCESS_KEY_ID": {Type: HardeningPolicyTypeDefault, Length: 20},
"LANGFUSE_S3_SECRET_ACCESS_KEY": {Type: HardeningPolicyTypeDefault, Length: 40},
"LANGFUSE_REDIS_AUTH": {Type: HardeningPolicyTypeHex, Length: 48},
"LANGFUSE_SALT": {Type: HardeningPolicyTypeHex, Length: 28},
"LANGFUSE_ENCRYPTION_KEY": {Type: HardeningPolicyTypeHex, Length: 64},
"LANGFUSE_NEXTAUTH_SECRET": {Type: HardeningPolicyTypeHex, Length: 32},
"LANGFUSE_INIT_PROJECT_PUBLIC_KEY": {Type: HardeningPolicyTypeUUID, Prefix: "pk-lf-"},
"LANGFUSE_INIT_PROJECT_SECRET_KEY": {Type: HardeningPolicyTypeUUID, Prefix: "sk-lf-"},
"LANGFUSE_AUTH_DISABLE_SIGNUP": {Type: HardeningPolicyTypeBoolTrue},
// LANGFUSE_PROJECT_ID, LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY are handled specially in syncLangfuseState
// LANGFUSE_INIT_USER_PASSWORD changes in web UI after first login, so we don't need to harden it
},
}
func DoHardening(s state.State, c checker.CheckResult) error {
var haveToCommit bool
installationID := system.GetInstallationID().String()
if id, _ := s.GetVar("INSTALLATION_ID"); id.Value != installationID {
if err := s.SetVar("INSTALLATION_ID", installationID); err != nil {
return fmt.Errorf("failed to set INSTALLATION_ID: %w", err)
}
haveToCommit = true
}
if licenseKey, exists := s.GetVar("LICENSE_KEY"); exists && licenseKey.Value != "" {
if info, err := sdk.IntrospectLicenseKey(licenseKey.Value); err != nil {
return fmt.Errorf("failed to introspect license key: %w", err)
} else if !info.IsValid() {
if err := s.SetVar("LICENSE_KEY", ""); err != nil {
return fmt.Errorf("failed to set LICENSE_KEY: %w", err)
}
haveToCommit = true
}
}
// harden langfuse vars only if neither containers nor volumes exist
// this prevents password changes when volumes with existing credentials are present
if vars, _ := s.GetVars(varsForHardening[HardeningAreaLangfuse]); !c.LangfuseInstalled && !c.LangfuseVolumesExist {
updateDefaultValues(vars)
if isChanged, err := replaceDefaultValues(s, vars, varsHardeningPolicies[HardeningAreaLangfuse]); err != nil {
return fmt.Errorf("failed to replace default values for langfuse: %w", err)
} else if isChanged {
haveToCommit = true
}
if isChanged, err := syncLangfuseState(s, vars); err != nil {
return fmt.Errorf("failed to sync langfuse vars: %w", err)
} else if isChanged {
haveToCommit = true
}
}
// harden graphiti vars only if neither containers nor volumes exist
// this prevents password changes when volumes with existing credentials are present
if vars, _ := s.GetVars(varsForHardening[HardeningAreaGraphiti]); !c.GraphitiInstalled && !c.GraphitiVolumesExist {
updateDefaultValues(vars)
if isChanged, err := replaceDefaultValues(s, vars, varsHardeningPolicies[HardeningAreaGraphiti]); err != nil {
return fmt.Errorf("failed to replace default values for graphiti: %w", err)
} else if isChanged {
haveToCommit = true
}
}
// harden pentagi vars only if neither containers nor volumes exist
// this prevents password changes when volumes with existing credentials are present
if vars, _ := s.GetVars(varsForHardening[HardeningAreaPentagi]); !c.PentagiInstalled && !c.PentagiVolumesExist {
updateDefaultValues(vars)
if isChanged, err := replaceDefaultValues(s, vars, varsHardeningPolicies[HardeningAreaPentagi]); err != nil {
return fmt.Errorf("failed to replace default values for pentagi: %w", err)
} else if isChanged {
haveToCommit = true
}
// sync scraper local URL access
if isChanged, err := syncScraperState(s, vars); err != nil {
return fmt.Errorf("failed to sync scraper state: %w", err)
} else if isChanged {
haveToCommit = true
}
}
if haveToCommit {
if err := s.Commit(); err != nil {
return fmt.Errorf("failed to commit vars: %w", err)
}
}
return nil
}
func syncValueToState(s state.State, curVar loader.EnvVar, newValue string) (loader.EnvVar, error) {
if err := s.SetVar(curVar.Name, newValue); err != nil {
return curVar, fmt.Errorf("failed to set var %s: %w", curVar.Name, err)
}
// get actual value from state and restore default value from previous step
newEnvVar, _ := s.GetVar(curVar.Name)
newEnvVar.Default = curVar.Value
return newEnvVar, nil
}
func syncScraperState(s state.State, vars map[string]loader.EnvVar) (bool, error) {
var isChanged bool
varName := "SCRAPER_PRIVATE_URL"
scraperPrivateURL, urlExists := vars[varName]
isDefaultScraperURL := urlExists && scraperPrivateURL.IsDefault()
scraperLocalUser, userExists := vars["LOCAL_SCRAPER_USERNAME"]
scraperLocalPassword, passwordExists := vars["LOCAL_SCRAPER_PASSWORD"]
isCredentialsExists := userExists && passwordExists
isCredentialsChanged := scraperLocalUser.IsChanged || scraperLocalPassword.IsChanged
if isDefaultScraperURL && isCredentialsExists && isCredentialsChanged {
parsedScraperPrivateURL, err := url.Parse(scraperPrivateURL.Value)
if err != nil {
return isChanged, fmt.Errorf("failed to parse scraper private URL: %w", err)
}
parsedScraperPrivateURL.User = url.UserPassword(scraperLocalUser.Value, scraperLocalPassword.Value)
syncedScraperPrivateURL, err := syncValueToState(s, scraperPrivateURL, parsedScraperPrivateURL.String())
if err != nil {
return isChanged, fmt.Errorf("failed to sync scraper private URL: %w", err)
}
vars[varName] = syncedScraperPrivateURL
if syncedScraperPrivateURL.IsChanged {
isChanged = true
}
}
return isChanged, nil
}
func syncLangfuseState(s state.State, vars map[string]loader.EnvVar) (bool, error) {
var isChanged bool
for varName, syncVarName := range varsHardeningSyncLangfuse {
envVar, exists := vars[varName]
if !exists {
continue
}
// don't change user values
if envVar.Value != "" {
continue
}
if syncVar, syncVarExists := vars[syncVarName]; syncVarExists {
syncedEnvVar, err := syncValueToState(s, envVar, syncVar.Value)
if err != nil {
return isChanged, fmt.Errorf("failed to sync var %s: %w", varName, err)
}
vars[varName] = syncedEnvVar
if syncedEnvVar.IsChanged {
isChanged = true
}
}
}
return isChanged, nil
}
func replaceDefaultValues(
s state.State, vars map[string]loader.EnvVar, policies map[string]HardeningPolicy,
) (bool, error) {
var (
err error
isChanged bool
)
for varName, envVar := range vars {
if policy, ok := policies[varName]; ok && envVar.IsDefault() {
envVar.Value, err = randomString(policy)
if err != nil {
return isChanged, fmt.Errorf("failed to generate random string for %s: %w", varName, err)
}
syncedEnvVar, err := syncValueToState(s, envVar, envVar.Value)
if err != nil {
return isChanged, fmt.Errorf("failed to sync var %s: %w", varName, err)
}
vars[varName] = syncedEnvVar
if syncedEnvVar.IsChanged {
isChanged = true
}
}
}
return isChanged, nil
}
func updateDefaultValues(vars map[string]loader.EnvVar) {
for varName, envVar := range vars {
if defVal, ok := varsForHardeningDefault[varName]; ok && envVar.Default == "" {
envVar.Default = defVal
vars[varName] = envVar
}
}
}
func randomString(policy HardeningPolicy) (string, error) {
switch policy.Type {
case HardeningPolicyTypeDefault:
return randStringAlpha(policy.Length)
case HardeningPolicyTypeHex:
return randStringHex(policy.Length)
case HardeningPolicyTypeUUID:
return randStringUUID(policy.Prefix)
case HardeningPolicyTypeBoolTrue:
return "true", nil
case HardeningPolicyTypeBoolFalse:
return "false", nil
default:
return "", fmt.Errorf("invalid hardening policy type: %s", policy.Type)
}
}
func randStringAlpha(length int) (string, error) {
bytes := make([]byte, length)
_, err := rand.Reader.Read(bytes)
if err != nil {
return "", err
}
charset := "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
for i, b := range bytes {
bytes[i] = charset[b%byte(len(charset))]
}
return string(bytes), nil
}
func randStringHex(length int) (string, error) {
byteLength := length/2 + 1
bytes := make([]byte, byteLength)
_, err := rand.Reader.Read(bytes)
if err != nil {
return "", err
}
hexString := hex.EncodeToString(bytes)
return hexString[:length], nil
}
func randStringUUID(prefix string) (string, error) {
return prefix + uuid.New().String(), nil
}
backend/cmd/installer/hardening/hardening_test.go
+1581
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/hardening/migrations.go
+88
View File
@@ -0,0 +1,88 @@
package hardening
import (
"os"
"slices"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/state"
"pentagi/cmd/installer/wizard/controller"
)
type checkPathType string
const (
directory checkPathType = "directory"
file checkPathType = "file"
)
func DoMigrateSettings(s state.State) error {
// migration from DOCKER_CERT_PATH to PENTAGI_DOCKER_CERT_PATH
dockerCertPathVar, exists := s.GetVar("DOCKER_CERT_PATH")
dockerCertPath := dockerCertPathVar.Value
if exists && dockerCertPath != "" {
exists = checkPathInHostFS(dockerCertPath, directory)
}
if exists && dockerCertPath != "" && dockerCertPath != controller.DefaultDockerCertPath {
if err := s.SetVar("PENTAGI_DOCKER_CERT_PATH", dockerCertPath); err != nil {
return err
}
if err := s.SetVar("DOCKER_CERT_PATH", controller.DefaultDockerCertPath); err != nil {
return err
}
}
configsPath := controller.GetEmbeddedLLMConfigsPath(files.NewFiles())
// migration from LLM_SERVER_CONFIG_PATH to PENTAGI_LLM_SERVER_CONFIG_PATH
llmServerConfigPathVar, exists := s.GetVar("LLM_SERVER_CONFIG_PATH")
llmServerConfigPath := llmServerConfigPathVar.Value
isEmbeddedCustomConfig := slices.Contains(configsPath, llmServerConfigPath) ||
llmServerConfigPath == controller.DefaultCustomConfigsPath
if exists && !isEmbeddedCustomConfig && llmServerConfigPath != "" {
exists = checkPathInHostFS(llmServerConfigPath, file)
}
if exists && !isEmbeddedCustomConfig && llmServerConfigPath != "" {
if err := s.SetVar("PENTAGI_LLM_SERVER_CONFIG_PATH", llmServerConfigPath); err != nil {
return err
}
if err := s.SetVar("LLM_SERVER_CONFIG_PATH", controller.DefaultCustomConfigsPath); err != nil {
return err
}
}
// migration from OLLAMA_SERVER_CONFIG_PATH to PENTAGI_OLLAMA_SERVER_CONFIG_PATH
ollamaServerConfigPathVar, exists := s.GetVar("OLLAMA_SERVER_CONFIG_PATH")
ollamaServerConfigPath := ollamaServerConfigPathVar.Value
isEmbeddedOllamaConfig := slices.Contains(configsPath, ollamaServerConfigPath) ||
ollamaServerConfigPath == controller.DefaultOllamaConfigsPath
if exists && !isEmbeddedOllamaConfig && ollamaServerConfigPath != "" {
exists = checkPathInHostFS(ollamaServerConfigPath, file)
}
if exists && !isEmbeddedOllamaConfig && ollamaServerConfigPath != "" {
if err := s.SetVar("PENTAGI_OLLAMA_SERVER_CONFIG_PATH", ollamaServerConfigPath); err != nil {
return err
}
if err := s.SetVar("OLLAMA_SERVER_CONFIG_PATH", controller.DefaultOllamaConfigsPath); err != nil {
return err
}
}
return nil
}
func checkPathInHostFS(path string, pathType checkPathType) bool {
info, err := os.Stat(path)
if err != nil {
return false
}
switch pathType {
case directory:
return info.IsDir()
case file:
return !info.IsDir()
default:
return false
}
}
backend/cmd/installer/hardening/migrations_test.go
+929
View File
@@ -0,0 +1,929 @@
package hardening
import (
"os"
"path/filepath"
"strings"
"testing"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/controller"
)
// Test 1: Successful migrations for all variables
func TestDoMigrateSettings_SuccessfulMigrations(t *testing.T) {
tests := []struct {
name string
setupFunc func(*testing.T) (string, func())
varName string
pentagiVarName string
defaultPath string
pathType checkPathType
customPath string
expectMigration bool
}{
{
name: "migrate DOCKER_CERT_PATH to PENTAGI_DOCKER_CERT_PATH",
setupFunc: func(t *testing.T) (string, func()) {
tmpDir, err := os.MkdirTemp("", "docker-certs-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
return tmpDir, func() { os.RemoveAll(tmpDir) }
},
varName: "DOCKER_CERT_PATH",
pentagiVarName: "PENTAGI_DOCKER_CERT_PATH",
defaultPath: controller.DefaultDockerCertPath,
pathType: directory,
expectMigration: true,
},
{
name: "migrate LLM_SERVER_CONFIG_PATH to PENTAGI_LLM_SERVER_CONFIG_PATH",
setupFunc: func(t *testing.T) (string, func()) {
tmpFile, err := os.CreateTemp("", "custom-*.yml")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
return tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
varName: "LLM_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
defaultPath: controller.DefaultCustomConfigsPath,
pathType: file,
expectMigration: true,
},
{
name: "migrate OLLAMA_SERVER_CONFIG_PATH to PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
setupFunc: func(t *testing.T) (string, func()) {
tmpFile, err := os.CreateTemp("", "ollama-*.yml")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
return tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
varName: "OLLAMA_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
defaultPath: controller.DefaultOllamaConfigsPath,
pathType: file,
expectMigration: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// setup temporary path
customPath, cleanup := tt.setupFunc(t)
defer cleanup()
// create mock state with custom path set
mockSt := &mockState{
vars: map[string]loader.EnvVar{
tt.varName: {
Name: tt.varName,
Value: customPath,
Line: 1,
IsChanged: false,
},
},
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify migration occurred
if tt.expectMigration {
// check that PENTAGI_* variable was set to custom path
pentagiVar, exists := mockSt.GetVar(tt.pentagiVarName)
if !exists {
t.Errorf("Expected %s to be set", tt.pentagiVarName)
} else if pentagiVar.Value != customPath {
t.Errorf("Expected %s = %q, got %q", tt.pentagiVarName, customPath, pentagiVar.Value)
}
// check that original variable was set to default path
originalVar, exists := mockSt.GetVar(tt.varName)
if !exists {
t.Errorf("Expected %s to be set", tt.varName)
} else if originalVar.Value != tt.defaultPath {
t.Errorf("Expected %s = %q, got %q", tt.varName, tt.defaultPath, originalVar.Value)
}
}
})
}
}
// Test 2: No migration when variable is not set
func TestDoMigrateSettings_VariableNotSet(t *testing.T) {
tests := []struct {
name string
pentagiVarName string
}{
{
name: "DOCKER_CERT_PATH not set",
pentagiVarName: "PENTAGI_DOCKER_CERT_PATH",
},
{
name: "LLM_SERVER_CONFIG_PATH not set",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
},
{
name: "OLLAMA_SERVER_CONFIG_PATH not set",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create mock state with no variables set
mockSt := &mockState{
vars: make(map[string]loader.EnvVar),
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify no migration occurred
_, exists := mockSt.GetVar(tt.pentagiVarName)
if exists {
t.Errorf("Expected %s to not be set", tt.pentagiVarName)
}
})
}
}
// Test 3: No migration when variable is empty
func TestDoMigrateSettings_EmptyVariable(t *testing.T) {
tests := []struct {
name string
varName string
pentagiVarName string
}{
{
name: "DOCKER_CERT_PATH is empty",
varName: "DOCKER_CERT_PATH",
pentagiVarName: "PENTAGI_DOCKER_CERT_PATH",
},
{
name: "LLM_SERVER_CONFIG_PATH is empty",
varName: "LLM_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
},
{
name: "OLLAMA_SERVER_CONFIG_PATH is empty",
varName: "OLLAMA_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create mock state with empty variable
mockSt := &mockState{
vars: map[string]loader.EnvVar{
tt.varName: {
Name: tt.varName,
Value: "",
Line: 1,
IsChanged: false,
},
},
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify no migration occurred
_, exists := mockSt.GetVar(tt.pentagiVarName)
if exists {
t.Errorf("Expected %s to not be set", tt.pentagiVarName)
}
})
}
}
// Test 4: No migration when path doesn't exist
func TestDoMigrateSettings_PathNotExist(t *testing.T) {
tests := []struct {
name string
varName string
pentagiVarName string
nonExistPath string
}{
{
name: "DOCKER_CERT_PATH points to non-existing directory",
varName: "DOCKER_CERT_PATH",
pentagiVarName: "PENTAGI_DOCKER_CERT_PATH",
nonExistPath: "/nonexistent/docker/certs",
},
{
name: "LLM_SERVER_CONFIG_PATH points to non-existing file",
varName: "LLM_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
nonExistPath: "/nonexistent/custom.provider.yml",
},
{
name: "OLLAMA_SERVER_CONFIG_PATH points to non-existing file",
varName: "OLLAMA_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
nonExistPath: "/nonexistent/ollama.provider.yml",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create mock state with non-existing path
mockSt := &mockState{
vars: map[string]loader.EnvVar{
tt.varName: {
Name: tt.varName,
Value: tt.nonExistPath,
Line: 1,
IsChanged: false,
},
},
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify no migration occurred
_, exists := mockSt.GetVar(tt.pentagiVarName)
if exists {
t.Errorf("Expected %s to not be set for non-existing path", tt.pentagiVarName)
}
})
}
}
// Test 5: No migration when variable already has default container path value
func TestDoMigrateSettings_AlreadyDefaultValue(t *testing.T) {
tests := []struct {
name string
varName string
pentagiVarName string
defaultPath string
description string
}{
{
name: "DOCKER_CERT_PATH already has default container path",
varName: "DOCKER_CERT_PATH",
pentagiVarName: "PENTAGI_DOCKER_CERT_PATH",
defaultPath: controller.DefaultDockerCertPath,
description: "Default container path should not be migrated",
},
{
name: "LLM_SERVER_CONFIG_PATH already has default container path",
varName: "LLM_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
defaultPath: controller.DefaultCustomConfigsPath,
description: "Default container path should not be migrated",
},
{
name: "OLLAMA_SERVER_CONFIG_PATH already has default container path",
varName: "OLLAMA_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
defaultPath: controller.DefaultOllamaConfigsPath,
description: "Default container path should not be migrated",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create mock state with default path
mockSt := &mockState{
vars: map[string]loader.EnvVar{
tt.varName: {
Name: tt.varName,
Value: tt.defaultPath,
Line: 1,
IsChanged: false,
},
},
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify no migration occurred
_, exists := mockSt.GetVar(tt.pentagiVarName)
if exists {
t.Errorf("Expected %s to not be set when already using default", tt.pentagiVarName)
}
// verify original variable was not changed
originalVar, exists := mockSt.GetVar(tt.varName)
if !exists {
t.Errorf("Expected %s to still exist", tt.varName)
} else if originalVar.Value != tt.defaultPath {
t.Errorf("Expected %s to remain %q, got %q", tt.varName, tt.defaultPath, originalVar.Value)
}
})
}
}
// Test 6: No migration for embedded LLM configs
func TestDoMigrateSettings_EmbeddedConfigs(t *testing.T) {
tests := []struct {
name string
varName string
pentagiVarName string
embeddedPath string
description string
}{
{
name: "LLM_SERVER_CONFIG_PATH with embedded config should not migrate",
varName: "LLM_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
embeddedPath: "/opt/pentagi/conf/llms/openai.yml",
description: "Embedded configs are inside docker image, no migration needed",
},
{
name: "OLLAMA_SERVER_CONFIG_PATH with embedded config should not migrate",
varName: "OLLAMA_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
embeddedPath: "/opt/pentagi/conf/llms/llama3.yml",
description: "Embedded configs are inside docker image, no migration needed",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create mock state with embedded config path
mockSt := &mockState{
vars: map[string]loader.EnvVar{
tt.varName: {
Name: tt.varName,
Value: tt.embeddedPath,
Line: 1,
IsChanged: false,
},
},
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify no migration occurred
pentagiVar, exists := mockSt.GetVar(tt.pentagiVarName)
if exists && pentagiVar.Value != "" {
t.Errorf("Expected %s to not be set for embedded config: %s", tt.pentagiVarName, tt.description)
}
// verify original variable was not changed
originalVar, exists := mockSt.GetVar(tt.varName)
if !exists {
t.Errorf("Expected %s to still exist", tt.varName)
} else if originalVar.Value != tt.embeddedPath {
t.Errorf("Expected %s to remain %q, got %q: %s", tt.varName, tt.embeddedPath, originalVar.Value, tt.description)
}
})
}
}
// Test 7: Wrong path type (file instead of directory and vice versa)
func TestDoMigrateSettings_WrongPathType(t *testing.T) {
tests := []struct {
name string
setupFunc func(*testing.T) (string, func())
varName string
pentagiVarName string
description string
}{
{
name: "DOCKER_CERT_PATH points to file instead of directory",
setupFunc: func(t *testing.T) (string, func()) {
tmpFile, err := os.CreateTemp("", "docker-cert-*")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
return tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
varName: "DOCKER_CERT_PATH",
pentagiVarName: "PENTAGI_DOCKER_CERT_PATH",
description: "File provided when directory expected",
},
{
name: "LLM_SERVER_CONFIG_PATH points to directory instead of file",
setupFunc: func(t *testing.T) (string, func()) {
tmpDir, err := os.MkdirTemp("", "llm-config-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
return tmpDir, func() { os.RemoveAll(tmpDir) }
},
varName: "LLM_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_LLM_SERVER_CONFIG_PATH",
description: "Directory provided when file expected",
},
{
name: "OLLAMA_SERVER_CONFIG_PATH points to directory instead of file",
setupFunc: func(t *testing.T) (string, func()) {
tmpDir, err := os.MkdirTemp("", "ollama-config-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
return tmpDir, func() { os.RemoveAll(tmpDir) }
},
varName: "OLLAMA_SERVER_CONFIG_PATH",
pentagiVarName: "PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
description: "Directory provided when file expected",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// setup wrong path type
wrongPath, cleanup := tt.setupFunc(t)
defer cleanup()
// create mock state with wrong path type
mockSt := &mockState{
vars: map[string]loader.EnvVar{
tt.varName: {
Name: tt.varName,
Value: wrongPath,
Line: 1,
IsChanged: false,
},
},
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify no migration occurred
_, exists := mockSt.GetVar(tt.pentagiVarName)
if exists {
t.Errorf("Expected %s to not be set for wrong path type: %s", tt.pentagiVarName, tt.description)
}
})
}
}
// Test 8: Error handling scenarios
func TestDoMigrateSettings_ErrorHandling(t *testing.T) {
tests := []struct {
name string
setupFunc func(*testing.T) (*mockStateWithErrors, string, func())
expectedError string
}{
{
name: "SetVar error for PENTAGI_DOCKER_CERT_PATH",
setupFunc: func(t *testing.T) (*mockStateWithErrors, string, func()) {
tmpDir, err := os.MkdirTemp("", "docker-certs-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
mockSt := &mockStateWithErrors{
vars: map[string]loader.EnvVar{
"DOCKER_CERT_PATH": {
Name: "DOCKER_CERT_PATH",
Value: tmpDir,
Line: 1,
},
},
setVarError: map[string]error{
"PENTAGI_DOCKER_CERT_PATH": mockError,
},
}
return mockSt, tmpDir, func() { os.RemoveAll(tmpDir) }
},
expectedError: "mocked error",
},
{
name: "SetVar error for DOCKER_CERT_PATH",
setupFunc: func(t *testing.T) (*mockStateWithErrors, string, func()) {
tmpDir, err := os.MkdirTemp("", "docker-certs-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
mockSt := &mockStateWithErrors{
vars: map[string]loader.EnvVar{
"DOCKER_CERT_PATH": {
Name: "DOCKER_CERT_PATH",
Value: tmpDir,
Line: 1,
},
},
setVarError: map[string]error{
"DOCKER_CERT_PATH": mockError,
},
}
return mockSt, tmpDir, func() { os.RemoveAll(tmpDir) }
},
expectedError: "mocked error",
},
{
name: "SetVar error for PENTAGI_LLM_SERVER_CONFIG_PATH",
setupFunc: func(t *testing.T) (*mockStateWithErrors, string, func()) {
tmpFile, err := os.CreateTemp("", "custom-*.yml")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
mockSt := &mockStateWithErrors{
vars: map[string]loader.EnvVar{
"LLM_SERVER_CONFIG_PATH": {
Name: "LLM_SERVER_CONFIG_PATH",
Value: tmpFile.Name(),
Line: 1,
},
},
setVarError: map[string]error{
"PENTAGI_LLM_SERVER_CONFIG_PATH": mockError,
},
}
return mockSt, tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
expectedError: "mocked error",
},
{
name: "SetVar error for PENTAGI_OLLAMA_SERVER_CONFIG_PATH",
setupFunc: func(t *testing.T) (*mockStateWithErrors, string, func()) {
tmpFile, err := os.CreateTemp("", "ollama-*.yml")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
mockSt := &mockStateWithErrors{
vars: map[string]loader.EnvVar{
"OLLAMA_SERVER_CONFIG_PATH": {
Name: "OLLAMA_SERVER_CONFIG_PATH",
Value: tmpFile.Name(),
Line: 1,
},
},
setVarError: map[string]error{
"PENTAGI_OLLAMA_SERVER_CONFIG_PATH": mockError,
},
}
return mockSt, tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
expectedError: "mocked error",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// setup mock state with error condition
mockSt, _, cleanup := tt.setupFunc(t)
defer cleanup()
// execute migration
err := DoMigrateSettings(mockSt)
// verify error was returned
if err == nil {
t.Error("Expected error but got none")
} else if err.Error() != tt.expectedError {
t.Errorf("Expected error %q, got %q", tt.expectedError, err.Error())
}
})
}
}
// Test 9: Combined migrations scenario
func TestDoMigrateSettings_CombinedMigrations(t *testing.T) {
tests := []struct {
name string
setupFunc func(*testing.T) (map[string]string, func())
expectedVars map[string]string
description string
}{
{
name: "migrate all three variables at once",
setupFunc: func(t *testing.T) (map[string]string, func()) {
// create temp directory for docker certs
dockerCertDir, err := os.MkdirTemp("", "docker-certs-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
// create temp file for LLM config
llmConfigFile, err := os.CreateTemp("", "custom-*.yml")
if err != nil {
os.RemoveAll(dockerCertDir)
t.Fatalf("Failed to create temp file: %v", err)
}
llmConfigFile.Close()
// create temp file for Ollama config
ollamaConfigFile, err := os.CreateTemp("", "ollama-*.yml")
if err != nil {
os.RemoveAll(dockerCertDir)
os.Remove(llmConfigFile.Name())
t.Fatalf("Failed to create temp file: %v", err)
}
ollamaConfigFile.Close()
paths := map[string]string{
"DOCKER_CERT_PATH": dockerCertDir,
"LLM_SERVER_CONFIG_PATH": llmConfigFile.Name(),
"OLLAMA_SERVER_CONFIG_PATH": ollamaConfigFile.Name(),
}
cleanup := func() {
os.RemoveAll(dockerCertDir)
os.Remove(llmConfigFile.Name())
os.Remove(ollamaConfigFile.Name())
}
return paths, cleanup
},
expectedVars: map[string]string{
"DOCKER_CERT_PATH": controller.DefaultDockerCertPath,
"LLM_SERVER_CONFIG_PATH": controller.DefaultCustomConfigsPath,
"OLLAMA_SERVER_CONFIG_PATH": controller.DefaultOllamaConfigsPath,
// PENTAGI_* vars will be checked separately as they contain dynamic temp paths
},
description: "All three migrations should complete successfully",
},
{
name: "migrate only DOCKER_CERT_PATH, others are default",
setupFunc: func(t *testing.T) (map[string]string, func()) {
dockerCertDir, err := os.MkdirTemp("", "docker-certs-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
paths := map[string]string{
"DOCKER_CERT_PATH": dockerCertDir,
"LLM_SERVER_CONFIG_PATH": controller.DefaultCustomConfigsPath,
"OLLAMA_SERVER_CONFIG_PATH": controller.DefaultOllamaConfigsPath,
}
cleanup := func() {
os.RemoveAll(dockerCertDir)
}
return paths, cleanup
},
expectedVars: map[string]string{
"DOCKER_CERT_PATH": controller.DefaultDockerCertPath,
"LLM_SERVER_CONFIG_PATH": controller.DefaultCustomConfigsPath,
"OLLAMA_SERVER_CONFIG_PATH": controller.DefaultOllamaConfigsPath,
},
description: "Only DOCKER_CERT_PATH should be migrated",
},
{
name: "migrate only config paths, DOCKER_CERT_PATH is default",
setupFunc: func(t *testing.T) (map[string]string, func()) {
// create temp file for LLM config
llmConfigFile, err := os.CreateTemp("", "custom-*.yml")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
llmConfigFile.Close()
// create temp file for Ollama config
ollamaConfigFile, err := os.CreateTemp("", "ollama-*.yml")
if err != nil {
os.Remove(llmConfigFile.Name())
t.Fatalf("Failed to create temp file: %v", err)
}
ollamaConfigFile.Close()
paths := map[string]string{
"DOCKER_CERT_PATH": controller.DefaultDockerCertPath,
"LLM_SERVER_CONFIG_PATH": llmConfigFile.Name(),
"OLLAMA_SERVER_CONFIG_PATH": ollamaConfigFile.Name(),
}
cleanup := func() {
os.Remove(llmConfigFile.Name())
os.Remove(ollamaConfigFile.Name())
}
return paths, cleanup
},
expectedVars: map[string]string{
"DOCKER_CERT_PATH": controller.DefaultDockerCertPath,
"LLM_SERVER_CONFIG_PATH": controller.DefaultCustomConfigsPath,
"OLLAMA_SERVER_CONFIG_PATH": controller.DefaultOllamaConfigsPath,
},
description: "Only config paths should be migrated",
},
{
name: "no migration for embedded configs",
setupFunc: func(t *testing.T) (map[string]string, func()) {
// create temp directory for docker certs
dockerCertDir, err := os.MkdirTemp("", "docker-certs-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
paths := map[string]string{
"DOCKER_CERT_PATH": dockerCertDir,
"LLM_SERVER_CONFIG_PATH": "/opt/pentagi/conf/llms/openai.yml", // embedded config
"OLLAMA_SERVER_CONFIG_PATH": "/opt/pentagi/conf/llms/llama3.yml", // embedded config
}
cleanup := func() {
os.RemoveAll(dockerCertDir)
}
return paths, cleanup
},
expectedVars: map[string]string{
"DOCKER_CERT_PATH": controller.DefaultDockerCertPath,
"LLM_SERVER_CONFIG_PATH": "/opt/pentagi/conf/llms/openai.yml", // should not change
"OLLAMA_SERVER_CONFIG_PATH": "/opt/pentagi/conf/llms/llama3.yml", // should not change
},
description: "Embedded configs should not be migrated, only DOCKER_CERT_PATH",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// setup paths and mock state
paths, cleanup := tt.setupFunc(t)
defer cleanup()
mockSt := &mockState{
vars: make(map[string]loader.EnvVar),
}
// populate mock state with initial values
for varName, varValue := range paths {
mockSt.vars[varName] = loader.EnvVar{
Name: varName,
Value: varValue,
Line: 1,
IsChanged: false,
}
}
// execute migration
err := DoMigrateSettings(mockSt)
if err != nil {
t.Fatalf("DoMigrateSettings() unexpected error = %v", err)
}
// verify all expected variables
for varName, expectedValue := range tt.expectedVars {
actualVar, exists := mockSt.GetVar(varName)
if !exists {
t.Errorf("Expected %s to be set", varName)
} else if actualVar.Value != expectedValue {
t.Errorf("Expected %s = %q, got %q", varName, expectedValue, actualVar.Value)
}
}
// verify PENTAGI_* variables were set correctly for non-default and non-embedded values
for varName, originalValue := range paths {
pentagiVarName := ""
defaultValue := ""
isEmbedded := false
switch varName {
case "DOCKER_CERT_PATH":
pentagiVarName = "PENTAGI_DOCKER_CERT_PATH"
defaultValue = controller.DefaultDockerCertPath
case "LLM_SERVER_CONFIG_PATH":
pentagiVarName = "PENTAGI_LLM_SERVER_CONFIG_PATH"
defaultValue = controller.DefaultCustomConfigsPath
// check if it's an embedded config path
isEmbedded = strings.HasPrefix(originalValue, "/opt/pentagi/conf/llms/")
case "OLLAMA_SERVER_CONFIG_PATH":
pentagiVarName = "PENTAGI_OLLAMA_SERVER_CONFIG_PATH"
defaultValue = controller.DefaultOllamaConfigsPath
// check if it's an embedded config path
isEmbedded = strings.HasPrefix(originalValue, "/opt/pentagi/conf/llms/")
}
// migration should only occur for non-default, non-embedded, existing files
shouldMigrate := originalValue != defaultValue && !isEmbedded
if shouldMigrate {
// check if file exists on host (migration only happens for existing files)
_, err := os.Stat(originalValue)
if err == nil {
// migration should have occurred
pentagiVar, exists := mockSt.GetVar(pentagiVarName)
if !exists {
t.Errorf("Expected %s to be set for non-default value", pentagiVarName)
} else if pentagiVar.Value != originalValue {
t.Errorf("Expected %s = %q, got %q", pentagiVarName, originalValue, pentagiVar.Value)
}
}
} else {
// migration should not have occurred
pentagiVar, exists := mockSt.GetVar(pentagiVarName)
if exists && pentagiVar.Value != "" {
t.Errorf("Expected %s to not be set for default/embedded value, but got %q", pentagiVarName, pentagiVar.Value)
}
}
}
})
}
}
// Test 10: checkPathInHostFS function
func TestCheckPathInHostFS(t *testing.T) {
tests := []struct {
name string
setupFunc func(*testing.T) (string, func())
pathType checkPathType
expectTrue bool
}{
{
name: "valid directory returns true for directory type",
setupFunc: func(t *testing.T) (string, func()) {
tmpDir, err := os.MkdirTemp("", "test-dir-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
return tmpDir, func() { os.RemoveAll(tmpDir) }
},
pathType: directory,
expectTrue: true,
},
{
name: "valid file returns false for directory type",
setupFunc: func(t *testing.T) (string, func()) {
tmpFile, err := os.CreateTemp("", "test-file-*")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
return tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
pathType: directory,
expectTrue: false,
},
{
name: "valid file returns true for file type",
setupFunc: func(t *testing.T) (string, func()) {
tmpFile, err := os.CreateTemp("", "test-file-*")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
tmpFile.Close()
return tmpFile.Name(), func() { os.Remove(tmpFile.Name()) }
},
pathType: file,
expectTrue: true,
},
{
name: "valid directory returns false for file type",
setupFunc: func(t *testing.T) (string, func()) {
tmpDir, err := os.MkdirTemp("", "test-dir-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
return tmpDir, func() { os.RemoveAll(tmpDir) }
},
pathType: file,
expectTrue: false,
},
{
name: "non-existent path returns false",
setupFunc: func(t *testing.T) (string, func()) {
tmpDir, err := os.MkdirTemp("", "test-dir-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
nonExistPath := filepath.Join(tmpDir, "nonexistent")
return nonExistPath, func() { os.RemoveAll(tmpDir) }
},
pathType: directory,
expectTrue: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
path, cleanup := tt.setupFunc(t)
defer cleanup()
result := checkPathInHostFS(path, tt.pathType)
if result != tt.expectTrue {
t.Errorf("checkPathInHostFS(%q, %v) = %v, want %v", path, tt.pathType, result, tt.expectTrue)
}
})
}
}
backend/cmd/installer/hardening/network.go
+65
View File
@@ -0,0 +1,65 @@
package hardening
import (
"os"
"pentagi/cmd/installer/state"
"pentagi/cmd/installer/wizard/controller"
)
func DoSyncNetworkSettings(s state.State) error {
// sync HTTP_PROXY or HTTPS_PROXY to PROXY_URL if they are set in the OS
httpProxy, httpProxyExists := os.LookupEnv("HTTP_PROXY")
if httpProxyExists && httpProxy != "" {
if err := s.SetVar("PROXY_URL", httpProxy); err != nil {
return err
}
}
httpsProxy, httpsProxyExists := os.LookupEnv("HTTPS_PROXY")
if httpsProxyExists && httpsProxy != "" {
if err := s.SetVar("PROXY_URL", httpsProxy); err != nil {
return err
}
}
dockerEnvVarsNames := []string{
"DOCKER_HOST",
"DOCKER_TLS_VERIFY",
"DOCKER_CERT_PATH",
"PENTAGI_DOCKER_CERT_PATH",
}
vars, exists := s.GetVars(dockerEnvVarsNames)
for _, envVar := range dockerEnvVarsNames {
if exists[envVar] && vars[envVar].Value != "" {
return nil // redefine is allowed only for unset docker connection settings
}
}
// get the environment variables from the OS
isOSDockerEnvVarsSet := false
osDockerEnvVars := make(map[string]string, len(dockerEnvVarsNames))
for _, envVar := range dockerEnvVarsNames {
value, exists := os.LookupEnv(envVar)
osDockerEnvVars[envVar] = value // set even empty value to avoid inconsistency while setting vars
if exists && value != "" {
isOSDockerEnvVarsSet = true
}
}
// do nothing if the OS docker environment variables are not set (use defaults)
if !isOSDockerEnvVarsSet {
return nil
}
// sync DOCKER_CERT_PATH to PENTAGI_DOCKER_CERT_PATH if it is set in the OS
dockerCertPath := osDockerEnvVars["DOCKER_CERT_PATH"]
if dockerCertPath != "" && checkPathInHostFS(dockerCertPath, directory) {
osDockerEnvVars["DOCKER_CERT_PATH"] = controller.DefaultDockerCertPath
osDockerEnvVars["PENTAGI_DOCKER_CERT_PATH"] = dockerCertPath
}
// sync all variables in the state at the same time to avoid inconsistencies
return s.SetVars(osDockerEnvVars)
}
backend/cmd/installer/hardening/network_test.go
+1051
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/loader/example_test.go
+73
View File
@@ -0,0 +1,73 @@
package loader
import (
"fmt"
"os"
"path/filepath"
)
// Example demonstrates the full workflow of loading, modifying, and saving .env files
func ExampleEnvFile_workflow() {
// Create a temporary .env file
tmpDir, _ := os.MkdirTemp("", "example")
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
initialContent := `# PentAGI Configuration
DATABASE_URL=postgres://localhost:5432/db
DEBUG=false
# API Settings
API_KEY=old_key`
os.WriteFile(envPath, []byte(initialContent), 0644)
// Step 1: Load existing .env file
envFile, err := LoadEnvFile(envPath)
if err != nil {
panic(err)
}
// Step 2: Display current values and defaults (only variables from file)
fmt.Println("Current configuration:")
fileVars := []string{"DATABASE_URL", "DEBUG", "API_KEY"}
for _, name := range fileVars {
envVar, exists := envFile.Get(name)
if !exists {
fmt.Printf("%s = (not present)\n", name)
continue
}
if envVar.IsPresent() && !envVar.IsComment {
fmt.Printf("%s = %s", name, envVar.Value)
if envVar.Default != "" && envVar.Default != envVar.Value {
fmt.Printf(" (default: %s)", envVar.Default)
}
if envVar.IsChanged {
fmt.Printf(" [modified]")
}
fmt.Println()
}
}
// Step 3: User modifies values
envFile.Set("DEBUG", "true")
envFile.Set("API_KEY", "new_secret_key")
envFile.Set("NEW_SETTING", "added_value")
// Step 4: Save changes (creates backup automatically)
err = envFile.Save(envPath)
if err != nil {
panic(err)
}
fmt.Println("\nConfiguration saved successfully!")
fmt.Println("Backup created in .bak directory")
// Output:
// Current configuration:
// DATABASE_URL = postgres://localhost:5432/db (default: postgres://pentagiuser:pentagipass@pgvector:5432/pentagidb?sslmode=disable)
// DEBUG = false
// API_KEY = old_key
//
// Configuration saved successfully!
// Backup created in .bak directory
}
backend/cmd/installer/loader/file.go
+243
View File
@@ -0,0 +1,243 @@
package loader
import (
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"time"
)
type EnvVar struct {
Name string // variable name
Value string // variable value
IsChanged bool // was the value changed manually
IsComment bool // is this line a comment (not saved, updated on value change)
Default string // default value from config struct (not saved, used for display)
Line int // line number in file (-1 if not present, e.g. for new vars)
}
func (e *EnvVar) IsDefault() bool {
return e.Value == e.Default || (e.Value == "" && e.Default != "")
}
func (e *EnvVar) IsPresent() bool {
return e.Line != -1
}
type EnvFile interface {
Del(name string)
Set(name, value string)
Get(name string) (EnvVar, bool)
GetAll() map[string]EnvVar
SetAll(vars map[string]EnvVar)
Save(path string) error
Clone() EnvFile
}
type envFile struct {
vars map[string]*EnvVar
perm os.FileMode
raw string
mx *sync.Mutex
}
func (e *envFile) Del(name string) {
e.mx.Lock()
defer e.mx.Unlock()
delete(e.vars, name)
}
func (e *envFile) Set(name, value string) {
e.mx.Lock()
defer e.mx.Unlock()
name, value = trim(name), trim(value)
if envVar, ok := e.vars[name]; !ok {
e.vars[name] = &EnvVar{
Name: name,
Value: value,
IsChanged: true,
Line: -1,
}
} else {
if envVar.Value != value {
envVar.IsChanged = true
envVar.Value = value
}
}
}
func (e *envFile) Get(name string) (EnvVar, bool) {
e.mx.Lock()
defer e.mx.Unlock()
if envVar, ok := e.vars[name]; !ok {
return EnvVar{
Name: name,
Line: -1,
}, false
} else {
return *envVar, true
}
}
func (e *envFile) GetAll() map[string]EnvVar {
e.mx.Lock()
defer e.mx.Unlock()
result := make(map[string]EnvVar, len(e.vars))
for name, envVar := range e.vars {
result[name] = *envVar
}
return result
}
func (e *envFile) SetAll(vars map[string]EnvVar) {
e.mx.Lock()
defer e.mx.Unlock()
for name := range vars {
envVar := vars[name]
e.vars[name] = &envVar
}
}
func (e *envFile) Save(path string) error {
e.mx.Lock()
defer e.mx.Unlock()
// check if there are any changes to the file to avoid unnecessary writes
curRaw := e.raw
e.patchRaw()
isChanged := e.raw != curRaw
for _, envVar := range e.vars {
if envVar.IsChanged {
isChanged = true
break
}
}
if !isChanged {
return nil
}
backupDir := filepath.Join(filepath.Dir(path), ".bak")
if err := os.MkdirAll(backupDir, 0755); err != nil {
return fmt.Errorf("failed to create backup directory: %w", err)
}
info, err := os.Stat(path)
if err == nil && info.IsDir() {
return fmt.Errorf("'%s' is a directory", path)
} else if err == nil {
curTimeStr := time.Unix(time.Now().Unix(), 0).Format("20060102150405")
backupPath := filepath.Join(backupDir, fmt.Sprintf("%s.%s", filepath.Base(path), curTimeStr))
if err := os.Rename(path, backupPath); err != nil {
return fmt.Errorf("failed to create backup file: %w", err)
}
}
if err := os.WriteFile(path, []byte(e.raw), e.perm); err != nil {
return fmt.Errorf("failed to write new file state: %w", err)
}
for _, envVar := range e.vars {
envVar.IsChanged = false
}
return nil
}
func (e *envFile) Clone() EnvFile {
e.mx.Lock()
defer e.mx.Unlock()
clone := envFile{
vars: make(map[string]*EnvVar, len(e.vars)),
perm: e.perm,
raw: e.raw,
mx: &sync.Mutex{},
}
for name, envVar := range e.vars {
v := *envVar
clone.vars[name] = &v
}
return &clone
}
func (e *envFile) patchRaw() {
lines := strings.Split(e.raw, "\n")
hasLastEmpty := len(lines) > 0 && trim(lines[len(lines)-1]) == ""
for ldx := len(lines) - 1; ldx >= 0 && trim(lines[ldx]) == ""; ldx-- {
lines = lines[:ldx]
}
// First pass: mark lines for deletion and update existing variables
var linesToDelete []int
for ldx, line := range lines {
line = trim(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
parts := strings.SplitN(line, "=", 2)
if len(parts) != 2 {
continue
}
varName := trim(parts[0])
// Check if this variable still exists
if envVar, exists := e.vars[varName]; exists {
if envVar.IsChanged && !envVar.IsComment {
lines[ldx] = fmt.Sprintf("%s=%s", envVar.Name, envVar.Value)
envVar.Line = ldx
}
} else {
// Mark line for deletion
linesToDelete = append(linesToDelete, ldx)
}
}
// Remove lines in reverse order to maintain indices
for i := len(linesToDelete) - 1; i >= 0; i-- {
lineIdx := linesToDelete[i]
lines = append(lines[:lineIdx], lines[lineIdx+1:]...)
// Update line numbers for remaining variables
for _, envVar := range e.vars {
if envVar.Line > lineIdx {
envVar.Line--
}
}
}
// Second pass: add new variables
for _, envVar := range e.vars {
if !envVar.IsChanged || envVar.IsComment {
continue
}
line := fmt.Sprintf("%s=%s", envVar.Name, envVar.Value)
if !envVar.IsPresent() || envVar.Line >= len(lines) {
lines = append(lines, line)
envVar.Line = len(lines) - 1
} else {
lines[envVar.Line] = line
}
}
if hasLastEmpty {
lines = append(lines, "")
}
e.raw = strings.Join(lines, "\n")
}
func trim(value string) string {
return strings.Trim(value, "\n\r\t ")
}
backend/cmd/installer/loader/loader.go
+131
View File
@@ -0,0 +1,131 @@
package loader
import (
"fmt"
"net/url"
"os"
"reflect"
"strconv"
"strings"
"sync"
"pentagi/pkg/config"
"github.com/caarlos0/env/v10"
)
func LoadEnvFile(path string) (EnvFile, error) {
info, err := os.Stat(path)
if err != nil {
return nil, fmt.Errorf("failed to stat '%s' file: %w", path, err)
} else if info.IsDir() {
return nil, fmt.Errorf("'%s' is a directory", path)
}
raw, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read '%s' file: %w", path, err)
}
envFile := &envFile{
vars: loadVars(string(raw)),
perm: info.Mode(),
raw: string(raw),
mx: &sync.Mutex{},
}
if err := setDefaultVars(envFile); err != nil {
return nil, fmt.Errorf("failed to set default vars: %w", err)
}
return envFile, nil
}
func loadVars(raw string) map[string]*EnvVar {
lines := strings.Split(string(raw), "\n")
vars := make(map[string]*EnvVar, len(lines))
for ldx, line := range lines {
envVar := &EnvVar{Line: ldx}
line = trim(line)
if line == "" {
continue
}
if strings.HasPrefix(line, "#") {
envVar.IsComment = true
line = trim(strings.TrimPrefix(line, "#"))
}
parts := strings.SplitN(line, "=", 2)
if len(parts) != 2 {
continue
}
envVar.Name = trim(parts[0])
envVar.Value = trim(stripComments(parts[1]))
envVar.IsChanged = envVar.Value != parts[1] || envVar.Name != parts[0]
if envVar.Name != "" {
vars[envVar.Name] = envVar
}
}
return vars
}
func stripComments(value string) string {
parts := strings.SplitN(value, " # ", 2)
if len(parts) == 2 {
return parts[0]
}
return value
}
func setDefaultVars(envFile *envFile) error {
var defaultConfig config.Config
if err := env.ParseWithOptions(&defaultConfig, env.Options{
FuncMap: map[reflect.Type]env.ParserFunc{
reflect.TypeOf(&url.URL{}): func(s string) (any, error) {
if s == "" {
return nil, nil
}
return url.Parse(s)
},
},
OnSet: func(tag string, value any, isDefault bool) {
if !isDefault {
return
}
var valueStr string
switch v := value.(type) {
case string:
valueStr = v
case *url.URL:
if v != nil {
valueStr = v.String()
}
case int:
valueStr = strconv.Itoa(v)
case bool:
valueStr = strconv.FormatBool(v)
default:
valueStr = fmt.Sprintf("%v", v)
}
if envVar, ok := envFile.vars[tag]; ok {
envVar.Default = valueStr
} else {
envFile.vars[tag] = &EnvVar{
Name: tag,
Value: "",
Default: valueStr,
Line: -1,
}
}
},
}); err != nil {
return fmt.Errorf("failed to parse env file: %w", err)
}
return nil
}
backend/cmd/installer/loader/loader_test.go
+698
View File
@@ -0,0 +1,698 @@
package loader
import (
"os"
"path/filepath"
"strings"
"sync"
"testing"
)
func containsLine(content, line string) bool {
lines := strings.Split(content, "\n")
for _, l := range lines {
if strings.TrimSpace(l) == line {
return true
}
}
return false
}
func TestLoadEnvFile(t *testing.T) {
tests := []struct {
name string
content string
wantErr bool
}{
{
name: "valid file",
content: `# Comment
VAR1=value1
VAR2=value2
# Another comment
VAR3=value3`,
wantErr: false,
},
{
name: "empty file",
content: "",
wantErr: false,
},
{
name: "comment only",
content: "# Just a comment",
wantErr: false,
},
{
name: "malformed lines",
content: `VAR1=value1
invalid line
VAR2=value2`,
wantErr: false,
},
{
name: "comments in value",
content: `VAR1=value1 # comment
VAR2=value2 # comment`,
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tmpFile := createTempFile(t, tt.content)
defer os.Remove(tmpFile)
envFile, err := LoadEnvFile(tmpFile)
if (err != nil) != tt.wantErr {
t.Errorf("LoadEnvFile() error = %v, wantErr %v", err, tt.wantErr)
return
}
if !tt.wantErr && envFile == nil {
t.Error("Expected envFile to be non-nil")
}
})
}
}
func TestLoadEnvFileErrors(t *testing.T) {
t.Run("non-existent file", func(t *testing.T) {
_, err := LoadEnvFile("/non/existent/file")
if err == nil {
t.Error("Expected error for non-existent file")
}
})
t.Run("directory instead of file", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
_, err := LoadEnvFile(tmpDir)
if err == nil {
t.Error("Expected error when path is directory")
}
})
}
func TestEnvVarMethods(t *testing.T) {
envVar := &EnvVar{
Name: "TEST_VAR",
Value: "test_value",
Default: "default_value",
Line: 5,
}
if envVar.IsDefault() {
t.Error("Expected IsDefault() to be false")
}
if !envVar.IsPresent() {
t.Error("Expected IsPresent() to be true")
}
envVar.Value = "default_value"
if !envVar.IsDefault() {
t.Error("Expected IsDefault() to be true")
}
envVar.Line = -1
if envVar.IsPresent() {
t.Error("Expected IsPresent() to be false")
}
}
func TestEnvFileSetGet(t *testing.T) {
envFile := &envFile{
vars: make(map[string]*EnvVar),
mx: &sync.Mutex{},
}
t.Run("set new variable", func(t *testing.T) {
envFile.Set("NEW_VAR", "new_value")
envVar, exists := envFile.Get("NEW_VAR")
if !exists {
t.Error("Expected NEW_VAR to exist")
}
if envVar.Name != "NEW_VAR" || envVar.Value != "new_value" {
t.Errorf("Expected NEW_VAR=new_value, got %s=%s", envVar.Name, envVar.Value)
}
if !envVar.IsChanged {
t.Error("Expected IsChanged to be true for new variable")
}
if envVar.Line != -1 {
t.Error("Expected Line to be -1 for new variable")
}
})
t.Run("update existing variable", func(t *testing.T) {
envFile.Set("NEW_VAR", "updated_value")
envVar, exists := envFile.Get("NEW_VAR")
if !exists {
t.Error("Expected NEW_VAR to exist")
}
if envVar.Value != "updated_value" {
t.Errorf("Expected updated_value, got %s", envVar.Value)
}
if !envVar.IsChanged {
t.Error("Expected IsChanged to remain true")
}
})
t.Run("set same value should not mark as changed", func(t *testing.T) {
// Reset IsChanged flag first
envFile.vars["NEW_VAR"].IsChanged = false
envFile.Set("NEW_VAR", "updated_value") // same value
envVar, exists := envFile.Get("NEW_VAR")
if !exists {
t.Error("Expected NEW_VAR to exist")
}
if envVar.IsChanged {
t.Error("Expected IsChanged to remain false when setting same value")
}
})
t.Run("get non-existent variable", func(t *testing.T) {
envVar, exists := envFile.Get("NON_EXISTENT")
if exists {
t.Error("Expected NON_EXISTENT to not exist")
}
if envVar.Name != "NON_EXISTENT" || envVar.Line != -1 {
t.Error("Expected empty EnvVar with Line=-1 for non-existent variable")
}
})
t.Run("trim whitespace", func(t *testing.T) {
envFile.Set(" TRIM_VAR ", " trim_value ")
envVar, exists := envFile.Get("TRIM_VAR")
if !exists {
t.Error("Expected TRIM_VAR to exist")
}
if envVar.Name != "TRIM_VAR" || envVar.Value != "trim_value" {
t.Errorf("Expected TRIM_VAR=trim_value, got %s=%s", envVar.Name, envVar.Value)
}
})
t.Run("delete variable", func(t *testing.T) {
envFile.Set("DELETE_VAR", "delete_value")
envVar, exists := envFile.Get("DELETE_VAR")
if !exists {
t.Error("Expected DELETE_VAR to exist before deletion")
}
if envVar.Value != "delete_value" {
t.Errorf("Expected DELETE_VAR value 'delete_value', got '%s'", envVar.Value)
}
envFile.Del("DELETE_VAR")
_, exists = envFile.Get("DELETE_VAR")
if exists {
t.Error("Expected DELETE_VAR to not exist after deletion")
}
})
t.Run("delete non-existent variable", func(t *testing.T) {
originalCount := len(envFile.GetAll())
envFile.Del("NON_EXISTENT_VAR")
if len(envFile.GetAll()) != originalCount {
t.Error("Deleting non-existent variable should not change variable count")
}
})
t.Run("get all variables", func(t *testing.T) {
allVars := envFile.GetAll()
if len(allVars) < 2 { // should have at least NEW_VAR and TRIM_VAR
t.Errorf("Expected at least 2 variables, got %d", len(allVars))
}
if newVar, exists := allVars["NEW_VAR"]; !exists {
t.Error("Expected NEW_VAR in GetAll result")
} else if newVar.Value != "updated_value" {
t.Errorf("Expected NEW_VAR value 'updated_value', got '%s'", newVar.Value)
}
})
t.Run("set all variables", func(t *testing.T) {
newVars := map[string]EnvVar{
"BATCH_VAR1": {Name: "BATCH_VAR1", Value: "batch_value1", IsChanged: true, Line: -1},
"BATCH_VAR2": {Name: "BATCH_VAR2", Value: "batch_value2", IsChanged: false, Line: 5},
"EXISTING_VAR": {Name: "EXISTING_VAR", Value: "overwritten", IsChanged: true, Line: 10},
}
envFile.SetAll(newVars)
// Check that all new variables were set
for name, expected := range newVars {
actual, exists := envFile.Get(name)
if !exists {
t.Errorf("Expected variable %s to exist after SetAll", name)
continue
}
if actual.Value != expected.Value {
t.Errorf("Variable %s: expected value %s, got %s", name, expected.Value, actual.Value)
}
if actual.IsChanged != expected.IsChanged {
t.Errorf("Variable %s: expected IsChanged %v, got %v", name, expected.IsChanged, actual.IsChanged)
}
if actual.Line != expected.Line {
t.Errorf("Variable %s: expected Line %d, got %d", name, expected.Line, actual.Line)
}
}
// Check that previous variables still exist
if _, exists := envFile.Get("NEW_VAR"); !exists {
t.Error("Expected NEW_VAR to still exist after SetAll")
}
})
t.Run("set all empty map", func(t *testing.T) {
originalCount := len(envFile.GetAll())
envFile.SetAll(map[string]EnvVar{})
if len(envFile.GetAll()) != originalCount {
t.Error("SetAll with empty map should not change existing variables")
}
})
}
func TestEnvFileSave(t *testing.T) {
content := `VAR1=value1
VAR2=value2`
tmpFile := createTempFile(t, content)
defer os.Remove(tmpFile)
envFile, err := LoadEnvFile(tmpFile)
if err != nil {
t.Fatalf("Failed to load env file: %v", err)
}
envFile.Set("VAR1", "new_value1")
envFile.Set("NEW_VAR", "new_value")
err = envFile.Save(tmpFile)
if err != nil {
t.Fatalf("Failed to save env file: %v", err)
}
// Check backup was created
backupDir := filepath.Join(filepath.Dir(tmpFile), ".bak")
entries, err := os.ReadDir(backupDir)
if err != nil {
t.Fatalf("Failed to read backup directory: %v", err)
}
if len(entries) == 0 {
t.Error("Expected backup file to be created")
}
// Check file content
savedContent, err := os.ReadFile(tmpFile)
if err != nil {
t.Fatalf("Failed to read saved file: %v", err)
}
expectedLines := []string{"VAR1=new_value1", "VAR2=value2", "NEW_VAR=new_value"}
savedLines := strings.Split(strings.TrimSpace(string(savedContent)), "\n")
for _, expected := range expectedLines {
found := false
for _, line := range savedLines {
if strings.TrimSpace(line) == expected {
found = true
break
}
}
if !found {
t.Errorf("Expected line '%s' not found in saved file", expected)
}
}
// Check IsChanged flags reset
for _, envVar := range envFile.GetAll() {
if envVar.IsChanged {
t.Errorf("Expected IsChanged to be false after save for %s", envVar.Name)
}
}
// Cleanup backup
os.RemoveAll(backupDir)
}
func TestEnvFileSaveNewFile(t *testing.T) {
envFile := &envFile{
vars: map[string]*EnvVar{
"VAR1": {Name: "VAR1", Value: "value1", IsChanged: true, Line: -1},
},
perm: 0644,
raw: "",
mx: &sync.Mutex{},
}
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
newFile := filepath.Join(tmpDir, "new.env")
err := envFile.Save(newFile)
if err != nil {
t.Fatalf("Failed to save new file: %v", err)
}
content, err := os.ReadFile(newFile)
if err != nil {
t.Fatalf("Failed to read new file: %v", err)
}
if !strings.Contains(string(content), "VAR1=value1") {
t.Error("Expected VAR1=value1 in new file")
}
}
func TestEnvFileSaveErrors(t *testing.T) {
const defaultEmptyContent = "# Empty file\n"
t.Run("save to directory", func(t *testing.T) {
envFile := &envFile{
vars: map[string]*EnvVar{
"VAR1": {Name: "VAR1", Value: "value1", IsChanged: true, Line: 0},
},
mx: &sync.Mutex{},
}
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
err := envFile.Save(tmpDir)
if err == nil {
t.Error("Expected error when saving to directory")
}
})
t.Run("save empty file", func(t *testing.T) {
envFile := &envFile{
vars: make(map[string]*EnvVar),
mx: &sync.Mutex{},
}
tmpFile := createTempFile(t, defaultEmptyContent)
defer os.Remove(tmpFile)
err := envFile.Save(tmpFile)
if err != nil {
t.Fatalf("Failed to save empty env file: %v", err)
}
content, err := os.ReadFile(tmpFile)
if err != nil {
t.Fatalf("Failed to read empty env file: %v", err)
}
if string(content) != defaultEmptyContent {
t.Errorf("Expected default empty content, got '%s'", string(content))
}
})
t.Run("save without changes", func(t *testing.T) {
envFile := &envFile{
vars: map[string]*EnvVar{
"VAR1": {Name: "VAR1", Value: "value1", IsChanged: false, Line: 0},
},
mx: &sync.Mutex{},
}
tmpFile := createTempFile(t, defaultEmptyContent)
defer os.Remove(tmpFile)
err := envFile.Save(tmpFile)
if err != nil {
t.Fatalf("Failed to save non changed env file: %v", err)
}
content, err := os.ReadFile(tmpFile)
if err != nil {
t.Fatalf("Failed to read empty env file: %v", err)
}
if string(content) != defaultEmptyContent {
t.Errorf("Expected default empty content, got '%s'", string(content))
}
})
}
func TestEnvFileClone(t *testing.T) {
original := &envFile{
vars: map[string]*EnvVar{
"VAR1": {Name: "VAR1", Value: "value1", IsChanged: true, Line: 0},
"VAR2": {Name: "VAR2", Value: "value2", IsChanged: false, Line: 1},
},
perm: 0644,
raw: "VAR1=value1\nVAR2=value2",
mx: &sync.Mutex{},
}
clone := original.Clone()
// Check independence
if clone == original {
t.Error("Clone should return different instance")
}
// Check content equality
if len(clone.GetAll()) != len(original.GetAll()) {
t.Error("Clone should have same number of variables")
}
for name, origVar := range original.GetAll() {
cloneVar, exists := clone.Get(name)
if !exists {
t.Errorf("Variable %s missing in clone", name)
continue
}
if cloneVar.Name != origVar.Name || cloneVar.Value != origVar.Value {
t.Errorf("Variable %s content mismatch in clone", name)
}
}
// Test modification independence
clone.Set("VAR1", "modified")
if original.vars["VAR1"].Value == "modified" {
t.Error("Modifying clone should not affect original")
}
}
func TestLoadVarsEdgeCases(t *testing.T) {
tests := []struct {
name string
content string
expected map[string]string
}{
{
name: "empty lines",
content: "\n\n\nVAR1=value1\n\n",
expected: map[string]string{"VAR1": "value1"},
},
{
name: "commented variables",
content: "#VAR1=commented\nVAR2=active",
expected: map[string]string{"VAR1": "commented", "VAR2": "active"},
},
{
name: "comments in value",
content: "VAR1=value1 # comment\nVAR2=value2 # comment",
expected: map[string]string{"VAR1": "value1", "VAR2": "value2"},
},
{
name: "variables with spaces",
content: "VAR1 = value1\n VAR2=value2 ",
expected: map[string]string{"VAR1": "value1", "VAR2": "value2"},
},
{
name: "variables with equals in value",
content: "VAR1=value=with=equals\nVAR2=url=https://example.com",
expected: map[string]string{"VAR1": "value=with=equals", "VAR2": "url=https://example.com"},
},
{
name: "invalid lines ignored",
content: "invalid line\nVAR1=value1\nanother invalid",
expected: map[string]string{"VAR1": "value1"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
vars := loadVars(tt.content)
for expectedName, expectedValue := range tt.expected {
envVar, exists := vars[expectedName]
if !exists {
t.Errorf("Expected variable %s not found", expectedName)
continue
}
if envVar.Value != expectedValue {
t.Errorf("Variable %s: expected value %s, got %s", expectedName, expectedValue, envVar.Value)
}
}
})
}
}
func TestPatchRaw(t *testing.T) {
envFile := &envFile{
vars: map[string]*EnvVar{
"EXISTING": {Name: "EXISTING", Value: "updated", IsChanged: true, Line: 1},
"NEW_VAR": {Name: "NEW_VAR", Value: "new_value", IsChanged: true, Line: -1},
},
raw: "# Comment line\nEXISTING=old_value\nUNCHANGED=unchanged\n",
mx: &sync.Mutex{},
}
envFile.patchRaw()
lines := strings.Split(envFile.raw, "\n")
// Check existing variable updated
if lines[1] != "EXISTING=updated" {
t.Errorf("Expected line 1 to be 'EXISTING=updated', got '%s'", lines[1])
}
// Check new variable added
found := false
for _, line := range lines {
if line == "NEW_VAR=new_value" {
found = true
break
}
}
if !found {
t.Error("Expected NEW_VAR=new_value to be added to file")
}
// Check comment not modified
if lines[0] != "# Comment line" {
t.Errorf("Expected comment line unchanged, got '%s'", lines[0])
}
// Check last line is empty
if lines[len(lines)-1] != "" {
t.Errorf("Expected last line to be empty, got '%s'", lines[len(lines)-1])
}
}
func TestEnvFileDelInSave(t *testing.T) {
content := `VAR1=value1
VAR2=value2
VAR3=value3`
tmpFile := createTempFile(t, content)
defer os.Remove(tmpFile)
envFile, err := LoadEnvFile(tmpFile)
if err != nil {
t.Fatalf("Failed to load env file: %v", err)
}
// Modify, add, and delete variables
envFile.Set("VAR1", "new_value1")
envFile.Set("NEW_VAR", "new_value")
envFile.Del("VAR2")
err = envFile.Save(tmpFile)
if err != nil {
t.Fatalf("Failed to save env file: %v", err)
}
// Check file content
savedContent, err := os.ReadFile(tmpFile)
if err != nil {
t.Fatalf("Failed to read saved file: %v", err)
}
contentStr := string(savedContent)
// Should contain updated and new variables
if !containsLine(contentStr, "VAR1=new_value1") {
t.Error("Expected VAR1 to be updated in saved file")
}
if !containsLine(contentStr, "NEW_VAR=new_value") {
t.Error("Expected NEW_VAR to be added to saved file")
}
if !containsLine(contentStr, "VAR3=value3") {
t.Error("Expected VAR3 to remain unchanged in saved file")
}
// Should not contain deleted variable
if containsLine(contentStr, "VAR2=value2") {
t.Error("Expected VAR2 to be removed from saved file")
}
// Cleanup backup
backupDir := filepath.Join(filepath.Dir(tmpFile), ".bak")
os.RemoveAll(backupDir)
}
func TestSetDefaultVarsNilURL(t *testing.T) {
envFile := &envFile{
vars: make(map[string]*EnvVar),
mx: &sync.Mutex{},
}
// This should not panic even with nil URL
err := setDefaultVars(envFile)
if err != nil {
t.Fatalf("setDefaultVars failed: %v", err)
}
// Check that STATIC_URL exists (it has envDefault empty, so should be nil URL)
if envVar, exists := envFile.vars["STATIC_URL"]; exists {
if envVar.Default != "" {
t.Errorf("Expected empty default for STATIC_URL, got '%s'", envVar.Default)
}
}
}
func TestSetDefaultVars(t *testing.T) {
envFile := &envFile{
vars: make(map[string]*EnvVar),
mx: &sync.Mutex{},
}
// This should not panic even with nil URL
err := setDefaultVars(envFile)
if err != nil {
t.Fatalf("setDefaultVars failed: %v", err)
}
// Check that all variables are not present and have default value the same as current value
for name, envVar := range envFile.vars {
if envVar.IsPresent() {
t.Errorf("Expected variable %s to be not present", name)
}
if !envVar.IsDefault() {
t.Errorf("Expected variable %s to have default value", name)
}
}
}
func createTempFile(t *testing.T, content string) string {
tmpFile, err := os.CreateTemp("", "test*.env")
if err != nil {
t.Fatalf("Failed to create temp file: %v", err)
}
if _, err := tmpFile.WriteString(content); err != nil {
t.Fatalf("Failed to write temp file: %v", err)
}
if err := tmpFile.Close(); err != nil {
t.Fatalf("Failed to close temp file: %v", err)
}
return tmpFile.Name()
}
backend/cmd/installer/main.go
+205
View File
@@ -0,0 +1,205 @@
package main
import (
"context"
"flag"
"fmt"
"log"
"os"
"os/signal"
"path/filepath"
"syscall"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/hardening"
"pentagi/cmd/installer/state"
"pentagi/cmd/installer/wizard"
"pentagi/pkg/version"
)
type Config struct {
envPath string
showVersion bool
}
func main() {
config := parseFlags(os.Args)
if config.showVersion {
fmt.Println(version.GetBinaryVersion())
os.Exit(0)
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
setupSignalHandler(cancel)
envPath, err := validateEnvPath(config.envPath)
if err != nil {
log.Fatalf("Error: %v", err)
}
appState, err := initializeState(envPath)
if err != nil {
log.Fatalf("Failed to initialize state: %v", err)
}
if err := hardening.DoMigrateSettings(appState); err != nil {
log.Fatalf("Failed to migrate settings: %v", err)
}
if err := hardening.DoSyncNetworkSettings(appState); err != nil {
log.Fatalf("Failed to sync network settings: %v", err)
}
checkResult, err := gatherSystemFacts(ctx, appState)
if err != nil {
log.Fatalf("Failed to gather system facts: %v", err)
}
printStartupInfo(envPath, checkResult)
if err := hardening.DoHardening(appState, checkResult); err != nil {
log.Fatalf("Failed to do hardening: %v", err)
}
if err := runApplication(ctx, appState, checkResult); err != nil {
log.Fatalf("Application error: %v", err)
}
cleanup(appState)
}
func parseFlags(args []string) Config {
var config Config
name := "installer"
if len(args) > 0 {
args, name = args[1:], filepath.Base(args[0])
}
flagSet := flag.NewFlagSet(name, flag.ContinueOnError)
flagSet.BoolVar(&config.showVersion, "v", false, "Show version information")
flagSet.StringVar(&config.envPath, "e", ".env", "Path to environment file")
flagSet.Usage = func() {
fmt.Fprintf(os.Stderr, "PentAGI Installer v%s\n\n", version.GetBinaryVersion())
fmt.Fprintf(os.Stderr, "Usage: %s [options]\n\n", name)
fmt.Fprintf(os.Stderr, "Options:\n")
flagSet.PrintDefaults()
fmt.Fprintf(os.Stderr, "\nExamples:\n")
fmt.Fprintf(os.Stderr, " %s # Use default .env file\n", name)
fmt.Fprintf(os.Stderr, " %s -e config/.env # Use custom env file\n", name)
fmt.Fprintf(os.Stderr, " %s -v # Show version\n", name)
}
flagSet.Parse(args)
return config
}
func setupSignalHandler(cancel context.CancelFunc) {
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
go func() {
sig := <-sigChan
log.Printf("Received signal: %v, initiating graceful shutdown...", sig)
cancel()
}()
}
func validateEnvPath(envPath string) (string, error) {
// convert to absolute path
absPath, err := filepath.Abs(envPath)
if err != nil {
return "", fmt.Errorf("invalid path '%s': %w", envPath, err)
}
// check if file exists
if info, err := os.Stat(absPath); os.IsNotExist(err) {
// file doesn't exist, check if we can create it in the directory
dir := filepath.Dir(absPath)
if _, err := os.Stat(dir); os.IsNotExist(err) {
if err := os.MkdirAll(dir, 0755); err != nil {
return "", fmt.Errorf("cannot create directory '%s': %w", dir, err)
}
} else if err != nil {
return "", fmt.Errorf("cannot access directory '%s': %w", dir, err)
}
// try to create initial env file
if err := createInitialEnvFile(absPath); err != nil {
return "", fmt.Errorf("cannot create env file '%s': %w", absPath, err)
}
} else if info.IsDir() {
return "", fmt.Errorf("'%s' is a directory", absPath)
} else if err != nil {
return "", fmt.Errorf("cannot access file '%s': %w", absPath, err)
}
return absPath, nil
}
func createInitialEnvFile(path string) error {
f := files.NewFiles()
content, err := f.GetContent(".env")
if err != nil {
return fmt.Errorf("cannot read .env file: %w", err)
}
content = fmt.Appendf(nil, `# PentAGI Environment Configuration
# Generated by PentAGI Installer v%s
#
# This file contains environment variables for PentAGI configuration.
# You can modify these values through the installer interface.
#
%s`, version.GetBinaryVersion(), string(content))
if err := os.WriteFile(path, content, 0600); err != nil {
return fmt.Errorf("cannot write .env file: %w", err)
}
return nil
}
func initializeState(envPath string) (state.State, error) {
appState, err := state.NewState(envPath)
if err != nil {
return nil, fmt.Errorf("failed to create state manager: %w", err)
}
return appState, nil
}
func gatherSystemFacts(ctx context.Context, appState state.State) (checker.CheckResult, error) {
result, err := checker.Gather(ctx, appState)
if err != nil {
return result, fmt.Errorf("failed to gather system facts: %w", err)
}
return result, nil
}
func printStartupInfo(envPath string, checkResult checker.CheckResult) {
fmt.Printf("PentAGI Installer v%s\n", version.GetBinaryVersion())
fmt.Printf("Environment file: %s\n", envPath)
if !checkResult.IsReadyToContinue() {
fmt.Println("⚠️ System is not ready to continue. Please resolve the issues above.")
} else {
fmt.Println("✅ System is ready to continue.")
}
}
func runApplication(ctx context.Context, appState state.State, checkResult checker.CheckResult) error {
return wizard.Run(ctx, appState, checkResult, files.NewFiles())
}
func cleanup(appState state.State) {
if appState.IsDirty() {
fmt.Println("You have pending changes.")
fmt.Println("Run the installer again to continue or commit your changes.")
}
}
backend/cmd/installer/main_test.go
+189
View File
@@ -0,0 +1,189 @@
package main
import (
"os"
"path/filepath"
"testing"
"pentagi/pkg/version"
)
func TestParseFlags(t *testing.T) {
tests := []struct {
name string
args []string
expectedEnv string
expectedVersion bool
}{
{
name: "default values",
args: []string{},
expectedEnv: ".env",
expectedVersion: false,
},
{
name: "custom env path",
args: []string{"-e", "config/.env"},
expectedEnv: "config/.env",
expectedVersion: false,
},
{
name: "version flag",
args: []string{"-v"},
expectedEnv: ".env",
expectedVersion: true,
},
{
name: "both flags",
args: []string{"-e", "test.env", "-v"},
expectedEnv: "test.env",
expectedVersion: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
config := parseFlags(append([]string{"test"}, tt.args...))
if config.envPath != tt.expectedEnv {
t.Errorf("Expected envPath %s, got %s", tt.expectedEnv, config.envPath)
}
if config.showVersion != tt.expectedVersion {
t.Errorf("Expected showVersion %v, got %v", tt.expectedVersion, config.showVersion)
}
})
}
}
func TestValidateEnvPath(t *testing.T) {
tmpDir := t.TempDir()
tests := []struct {
name string
path string
setup func() string
expectError bool
}{
{
name: "existing file",
setup: func() string {
path := filepath.Join(tmpDir, "existing.env")
os.WriteFile(path, []byte("VAR=value"), 0644)
return path
},
expectError: false,
},
{
name: "non-existent file in existing directory",
setup: func() string {
return filepath.Join(tmpDir, "new.env")
},
expectError: false,
},
{
name: "non-existent directory",
setup: func() string {
return filepath.Join(tmpDir, "nonexistent", "file.env")
},
expectError: false,
},
{
name: "directory instead of file",
setup: func() string {
os.Mkdir(filepath.Join(tmpDir, "dir"), 0755)
return filepath.Join(tmpDir, "dir")
},
expectError: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
path := tt.setup()
result, err := validateEnvPath(path)
if tt.expectError {
if err == nil {
t.Error("Expected error but got none")
}
} else {
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if result == "" {
t.Error("Expected non-empty result path")
}
// Check that file exists after validation
if _, err := os.Stat(result); os.IsNotExist(err) {
t.Error("Expected file to exist after validation")
}
}
})
}
}
func TestCreateEmptyEnvFile(t *testing.T) {
tmpDir := t.TempDir()
path := filepath.Join(tmpDir, "test.env")
err := createInitialEnvFile(path)
if err != nil {
t.Fatalf("Failed to create empty env file: %v", err)
}
// Check file exists
if _, err := os.Stat(path); os.IsNotExist(err) {
t.Error("Expected file to be created")
}
// Check file content
content, err := os.ReadFile(path)
if err != nil {
t.Fatalf("Failed to read created file: %v", err)
}
contentStr := string(content)
if !containsString(contentStr, "PentAGI Environment Configuration") {
t.Error("Expected file to contain header comment")
}
if !containsString(contentStr, version.GetBinaryVersion()) {
t.Error("Expected file to contain version")
}
}
func TestInitializeState(t *testing.T) {
tmpDir := t.TempDir()
envPath := filepath.Join(tmpDir, "test.env")
// Create test env file
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := initializeState(envPath)
if err != nil {
t.Fatalf("Failed to initialize state: %v", err)
}
if state == nil {
t.Error("Expected non-nil state")
}
// Test that state can access variables
envVar, exists := state.GetVar("VAR1")
if !exists {
t.Error("Expected VAR1 to exist in state")
}
if envVar.Value != "value1" {
t.Errorf("Expected VAR1 value 'value1', got '%s'", envVar.Value)
}
}
func containsString(s, substr string) bool {
return len(s) >= len(substr) &&
(s == substr ||
containsString(s[1:], substr) ||
(len(s) > 0 && s[:len(substr)] == substr))
}
backend/cmd/installer/navigator/navigator.go
+99
View File
@@ -0,0 +1,99 @@
package navigator
import (
"strings"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/state"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/models"
)
type Navigator interface {
Push(screenID models.ScreenID)
Pop() models.ScreenID
Current() models.ScreenID
CanGoBack() bool
GetStack() NavigatorStack
String() string
}
// navigator handles screen navigation and step persistence
type navigator struct {
stack NavigatorStack
stateManager state.State
}
type NavigatorStack []models.ScreenID
func (s NavigatorStack) Strings() []string {
strings := make([]string, len(s))
for i, v := range s {
strings[i] = string(v)
}
return strings
}
func (s NavigatorStack) String() string {
return strings.Join(s.Strings(), " -> ")
}
func NewNavigator(state state.State, checkResult checker.CheckResult) Navigator {
logger.Log("[Nav] NEW: %s", strings.Join(state.GetStack(), " -> "))
if !checkResult.IsReadyToContinue() {
state.SetStack([]string{string(models.WelcomeScreen)})
return &navigator{
stack: []models.ScreenID{models.WelcomeScreen},
stateManager: state,
}
}
stack := make([]models.ScreenID, 0)
for _, screenID := range state.GetStack() {
stack = append(stack, models.ScreenID(screenID))
}
return &navigator{
stack: stack,
stateManager: state,
}
}
func (n *navigator) Push(screenID models.ScreenID) {
logger.Log("[Nav] PUSH: %s -> %s", n.Current(), screenID)
n.stack = append(n.stack, screenID)
n.stateManager.SetStack(n.stack.Strings())
}
func (n *navigator) Pop() models.ScreenID {
current := n.Current()
if len(n.stack) <= 1 {
return current
}
n.stack = n.stack[:len(n.stack)-1]
previous := n.Current()
n.stateManager.SetStack(n.stack.Strings())
logger.Log("[Nav] POP: %s -> %s", current, previous)
return previous
}
func (n *navigator) Current() models.ScreenID {
if len(n.stack) == 0 {
return models.WelcomeScreen
}
return n.stack[len(n.stack)-1]
}
func (n *navigator) CanGoBack() bool {
return len(n.stack) > 1
}
func (n *navigator) GetStack() NavigatorStack {
return n.stack
}
func (n *navigator) String() string {
return n.stack.String()
}
backend/cmd/installer/navigator/navigator_test.go
+321
View File
@@ -0,0 +1,321 @@
package navigator
import (
"fmt"
"reflect"
"testing"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/models"
)
type mockState struct {
stack []string
}
func (m *mockState) Exists() bool { return true }
func (m *mockState) Reset() error { return nil }
func (m *mockState) Commit() error { return nil }
func (m *mockState) IsDirty() bool { return false }
func (m *mockState) GetEulaConsent() bool { return false }
func (m *mockState) SetEulaConsent() error { return nil }
func (m *mockState) GetStack() []string { return m.stack }
func (m *mockState) SetStack(stack []string) error { m.stack = stack; return nil }
func (m *mockState) GetVar(string) (loader.EnvVar, bool) { return loader.EnvVar{}, false }
func (m *mockState) SetVar(string, string) error { return nil }
func (m *mockState) ResetVar(string) error { return nil }
func (m *mockState) GetVars([]string) (map[string]loader.EnvVar, map[string]bool) {
return nil, nil
}
func (m *mockState) SetVars(map[string]string) error { return nil }
func (m *mockState) ResetVars([]string) error { return nil }
func (m *mockState) GetAllVars() map[string]loader.EnvVar { return nil }
func (m *mockState) GetEnvPath() string { return "" }
func newMockCheckResult() checker.CheckResult {
return checker.CheckResult{
EnvFileExists: true,
EnvDirWritable: true,
DockerApiAccessible: true,
WorkerEnvApiAccessible: true,
DockerComposeInstalled: true,
DockerVersionOK: true,
DockerComposeVersionOK: true,
SysNetworkOK: true,
SysCPUOK: true,
SysMemoryOK: true,
SysDiskFreeSpaceOK: true,
}
}
func newTestNavigator() (Navigator, *mockState) {
mockState := &mockState{}
nav := NewNavigator(mockState, newMockCheckResult())
return nav, mockState
}
func TestNewNavigator(t *testing.T) {
nav, state := newTestNavigator()
if nav.Current() != models.WelcomeScreen {
t.Errorf("expected WelcomeScreen, got %s", nav.Current())
}
if nav.CanGoBack() {
t.Error("new navigator should not allow going back")
}
if len(state.stack) != 0 {
t.Errorf("expected empty state stack, got %v", state.stack)
}
}
func TestPushSingleScreen(t *testing.T) {
nav, state := newTestNavigator()
nav.Push(models.MainMenuScreen)
if nav.Current() != models.MainMenuScreen {
t.Errorf("expected MainMenuScreen, got %s", nav.Current())
}
expected := []string{string(models.MainMenuScreen)}
if !reflect.DeepEqual(state.stack, expected) {
t.Errorf("expected state stack %v, got %v", expected, state.stack)
}
}
func TestPushMultipleScreens(t *testing.T) {
nav, state := newTestNavigator()
screens := []models.ScreenID{
models.MainMenuScreen,
models.LLMProvidersScreen,
models.LLMProviderOpenAIScreen,
}
for _, screen := range screens {
nav.Push(screen)
}
if nav.Current() != models.LLMProviderOpenAIScreen {
t.Errorf("expected LLMProviderFormScreen, got %s", nav.Current())
}
expected := []string{
string(models.MainMenuScreen),
string(models.LLMProvidersScreen),
string(models.LLMProviderOpenAIScreen),
}
if !reflect.DeepEqual(state.stack, expected) {
t.Errorf("expected state stack %v, got %v", expected, state.stack)
}
}
func TestCanGoBack(t *testing.T) {
nav, _ := newTestNavigator()
if nav.CanGoBack() {
t.Error("empty navigator should not allow going back")
}
nav.Push(models.MainMenuScreen)
if nav.CanGoBack() {
t.Error("single screen navigator should not allow going back")
}
nav.Push(models.LLMProvidersScreen)
if !nav.CanGoBack() {
t.Error("multi-screen navigator should allow going back")
}
}
func TestPopNormalCase(t *testing.T) {
nav, state := newTestNavigator()
nav.Push(models.MainMenuScreen)
nav.Push(models.LLMProvidersScreen)
nav.Push(models.LLMProviderOpenAIScreen)
previous := nav.Pop()
if previous != models.LLMProvidersScreen {
t.Errorf("expected LLMProvidersScreen, got %s", previous)
}
if nav.Current() != models.LLMProvidersScreen {
t.Errorf("expected current LLMProvidersScreen, got %s", nav.Current())
}
expected := []string{string(models.MainMenuScreen), string(models.LLMProvidersScreen)}
if !reflect.DeepEqual(state.stack, expected) {
t.Errorf("expected state stack %v, got %v", expected, state.stack)
}
}
func TestPopEmptyStack(t *testing.T) {
nav, _ := newTestNavigator()
result := nav.Pop()
if result != models.WelcomeScreen {
t.Errorf("expected WelcomeScreen, got %s", result)
}
if nav.Current() != models.WelcomeScreen {
t.Errorf("expected current WelcomeScreen, got %s", nav.Current())
}
}
func TestPopSingleScreen(t *testing.T) {
nav, _ := newTestNavigator()
nav.Push(models.MainMenuScreen)
result := nav.Pop()
if result != models.MainMenuScreen {
t.Errorf("expected MainMenuScreen, got %s", result)
}
if nav.Current() != models.MainMenuScreen {
t.Errorf("expected current MainMenuScreen, got %s", nav.Current())
}
}
func TestCurrentEmptyStack(t *testing.T) {
nav, _ := newTestNavigator()
if nav.Current() != models.WelcomeScreen {
t.Errorf("expected WelcomeScreen for empty stack, got %s", nav.Current())
}
}
func TestGetStack(t *testing.T) {
nav, _ := newTestNavigator()
screens := []models.ScreenID{
models.MainMenuScreen,
models.ToolsScreen,
models.DockerFormScreen,
}
for _, screen := range screens {
nav.Push(screen)
}
stack := nav.GetStack()
expected := NavigatorStack{models.MainMenuScreen, models.ToolsScreen, models.DockerFormScreen}
if !reflect.DeepEqual(stack, expected) {
t.Errorf("expected stack %v, got %v", expected, stack)
}
}
func TestNavigatorStackStrings(t *testing.T) {
stack := NavigatorStack{models.WelcomeScreen, models.MainMenuScreen, models.LLMProvidersScreen}
strings := stack.Strings()
expected := []string{
string(models.WelcomeScreen),
string(models.MainMenuScreen),
string(models.LLMProvidersScreen),
}
if !reflect.DeepEqual(strings, expected) {
t.Errorf("expected strings %v, got %v", expected, strings)
}
}
func TestNavigatorStackString(t *testing.T) {
stack := NavigatorStack{models.WelcomeScreen, models.MainMenuScreen, models.LLMProvidersScreen}
result := stack.String()
expected := fmt.Sprintf("%s -> %s -> %s", models.WelcomeScreen, models.MainMenuScreen, models.LLMProvidersScreen)
if result != expected {
t.Errorf("expected string %q, got %q", expected, result)
}
}
func TestNavigatorStackStringEmpty(t *testing.T) {
stack := NavigatorStack{}
result := stack.String()
expected := ""
if result != expected {
t.Errorf("expected empty string, got %q", result)
}
}
func TestNavigatorString(t *testing.T) {
nav, _ := newTestNavigator()
nav.Push(models.MainMenuScreen)
nav.Push(models.ToolsScreen)
result := nav.String()
expected := string(models.MainMenuScreen) + " -> " + string(models.ToolsScreen)
if result != expected {
t.Errorf("expected string %q, got %q", expected, result)
}
}
func TestComplexNavigationFlow(t *testing.T) {
nav, state := newTestNavigator()
// simulate typical navigation flow
nav.Push(models.MainMenuScreen)
nav.Push(models.LLMProvidersScreen)
nav.Push(models.LLMProviderOpenAIScreen)
// go back once
nav.Pop()
if nav.Current() != models.LLMProvidersScreen {
t.Errorf("expected LLMProvidersScreen after pop, got %s", nav.Current())
}
// navigate to different branch
nav.Push(models.ToolsScreen)
nav.Push(models.DockerFormScreen)
if nav.Current() != models.DockerFormScreen {
t.Errorf("expected DockerFormScreen, got %s", nav.Current())
}
// verify final state
expected := []string{
string(models.MainMenuScreen),
string(models.LLMProvidersScreen),
string(models.ToolsScreen),
string(models.DockerFormScreen),
}
if !reflect.DeepEqual(state.stack, expected) {
t.Errorf("expected final state %v, got %v", expected, state.stack)
}
if !nav.CanGoBack() {
t.Error("should be able to go back in complex flow")
}
}
func TestStateIntegrationWithExistingStack(t *testing.T) {
// test navigator initialization with existing stack
existingStack := []string{string(models.MainMenuScreen), string(models.ToolsScreen)}
mockState := &mockState{stack: existingStack}
nav := NewNavigator(mockState, newMockCheckResult())
// navigator should start with the last screen in the stack
if nav.Current() != models.ToolsScreen {
t.Errorf("expected ToolsScreen on new navigator, got %s", nav.Current())
}
if len(nav.GetStack()) != 2 {
t.Errorf("expected 2 screens in navigator stack, got %v", nav.GetStack())
}
}
backend/cmd/installer/processor/compose.go
+220
View File
@@ -0,0 +1,220 @@
package processor
import (
"context"
"fmt"
"os"
"os/exec"
"path/filepath"
"time"
)
const (
composeFilePentagi = "docker-compose.yml"
composeFileGraphiti = "docker-compose-graphiti.yml"
composeFileLangfuse = "docker-compose-langfuse.yml"
composeFileObservability = "docker-compose-observability.yml"
)
var composeOperationAllStacksOrder = map[ProcessorOperation][]ProductStack{
ProcessorOperationStart: {ProductStackObservability, ProductStackLangfuse, ProductStackGraphiti, ProductStackPentagi},
ProcessorOperationStop: {ProductStackPentagi, ProductStackGraphiti, ProductStackLangfuse, ProductStackObservability},
ProcessorOperationUpdate: {ProductStackObservability, ProductStackLangfuse, ProductStackGraphiti, ProductStackPentagi},
ProcessorOperationDownload: {ProductStackObservability, ProductStackLangfuse, ProductStackGraphiti, ProductStackPentagi},
ProcessorOperationRemove: {ProductStackObservability, ProductStackLangfuse, ProductStackGraphiti, ProductStackPentagi},
ProcessorOperationPurge: {ProductStackObservability, ProductStackLangfuse, ProductStackGraphiti, ProductStackPentagi},
}
type composeOperationsImpl struct {
processor *processor
}
func newComposeOperations(p *processor) composeOperations {
return &composeOperationsImpl{processor: p}
}
// startStack starts Docker Compose stack with dependency ordering
func (c *composeOperationsImpl) startStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationStart, "start")
}
// stopStack stops Docker Compose stack
func (c *composeOperationsImpl) stopStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationStop, "stop")
}
// restartStack restarts Docker Compose stack (stop + start to avoid race conditions for dependencies)
func (c *composeOperationsImpl) restartStack(ctx context.Context, stack ProductStack, state *operationState) error {
if err := c.stopStack(ctx, stack, state); err != nil {
return err
}
// brief pause to ensure clean shutdown
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(2 * time.Second):
}
return c.startStack(ctx, stack, state)
}
// updateStack performs rolling update with health checks (also used for install)
func (c *composeOperationsImpl) updateStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationUpdate, "up", "-d")
}
func (c *composeOperationsImpl) downloadStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationDownload, "pull")
}
func (c *composeOperationsImpl) removeStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationRemove, "down")
}
func (c *composeOperationsImpl) purgeStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationPurge, "down", "-v")
}
// purgeImagesStack is a stricter purge that also removes images referenced by the compose services
func (c *composeOperationsImpl) purgeImagesStack(ctx context.Context, stack ProductStack, state *operationState) error {
return c.performStackOperation(ctx, stack, state, ProcessorOperationPurge, "down", "--rmi", "all", "-v")
}
func (c *composeOperationsImpl) performStackOperation(
ctx context.Context, stack ProductStack, state *operationState, operation ProcessorOperation, args ...string,
) error {
switch stack {
case ProductStackPentagi:
// destructive operations cannot run without the compose file and have
// nothing to act on when the stack was never extracted
if isDestructiveComposeOperation(operation) && !c.composeFileExists(stack) {
return nil
}
return c.wrapPerformStackCommand(ctx, stack, state, operation, args...)
case ProductStackLangfuse, ProductStackObservability, ProductStackGraphiti:
switch operation {
// for destructive operations we must always allow compose to run, even if stack is disabled/external now
case ProcessorOperationRemove, ProcessorOperationPurge, ProcessorOperationStop:
// but only if the compose file exists; a missing file means the stack
// was never extracted and there is nothing for compose to remove
if !c.composeFileExists(stack) {
return nil
}
return c.wrapPerformStackCommand(ctx, stack, state, operation, args...)
// for non-destructive operations (start/update/download) honor embedded mode only
default:
if c.processor.isEmbeddedDeployment(stack) {
return c.wrapPerformStackCommand(ctx, stack, state, operation, args...)
}
return nil
}
case ProductStackAll, ProductStackCompose:
for _, s := range composeOperationAllStacksOrder[operation] {
if err := c.performStackOperation(ctx, s, state, operation, args...); err != nil {
return err
}
}
return nil
default:
return fmt.Errorf("operation %s not applicable for stack %s", operation, stack)
}
}
func (c *composeOperationsImpl) wrapPerformStackCommand(
ctx context.Context, stack ProductStack, state *operationState, operation ProcessorOperation, args ...string,
) error {
msgs := c.getMessages(stack, operation)
c.processor.appendLog(msgs.Enter, stack, state)
err := c.performStackCommand(ctx, stack, state, args...)
if err != nil {
c.processor.appendLog(fmt.Sprintf("%s: %s\n", msgs.Error, err.Error()), stack, state)
} else {
c.processor.appendLog(msgs.Exit+"\n", stack, state)
}
return err
}
func (c *composeOperationsImpl) performStackCommand(
ctx context.Context, stack ProductStack, state *operationState, args ...string,
) error {
envPath := c.processor.state.GetEnvPath()
composeFile, err := c.determineComposeFile(stack)
if err != nil {
return err
}
workingDir := filepath.Dir(envPath)
composePath := filepath.Join(workingDir, composeFile)
// check if files exist
if err := c.processor.isFileExists(composePath); err != nil {
return err
}
if err := c.processor.isFileExists(envPath); err != nil {
return err
}
// build docker compose command
args = append([]string{"compose", "--env-file", envPath, "-f", composePath}, args...)
cmd := exec.CommandContext(ctx, "docker", args...)
cmd.Dir = workingDir
cmd.Env = os.Environ()
// stacks are processed one by one, so we can ignore orphans
// orphans containers are removed by specific stack operations in main logic
cmd.Env = append(cmd.Env, "COMPOSE_IGNORE_ORPHANS=1")
// force Python unbuffered output to prevent incomplete data loss
cmd.Env = append(cmd.Env, "PYTHONUNBUFFERED=1")
return c.processor.runCommand(cmd, stack, state)
}
// isDestructiveComposeOperation reports whether the operation tears a stack down
// (and therefore may target a stack that is disabled/external/never extracted)
func isDestructiveComposeOperation(operation ProcessorOperation) bool {
switch operation {
case ProcessorOperationRemove, ProcessorOperationPurge, ProcessorOperationStop:
return true
default:
return false
}
}
// composeFileExists reports whether the stack's compose file is present on disk;
// destructive operations cannot run (and have nothing to remove) without it
func (c *composeOperationsImpl) composeFileExists(stack ProductStack) bool {
composeFile, err := c.determineComposeFile(stack)
if err != nil {
return false
}
composePath := filepath.Join(filepath.Dir(c.processor.state.GetEnvPath()), composeFile)
return c.processor.isFileExists(composePath) == nil
}
func (c *composeOperationsImpl) determineComposeFile(stack ProductStack) (string, error) {
switch stack {
case ProductStackPentagi:
return composeFilePentagi, nil
case ProductStackGraphiti:
return composeFileGraphiti, nil
case ProductStackLangfuse:
return composeFileLangfuse, nil
case ProductStackObservability:
return composeFileObservability, nil
default:
return "", fmt.Errorf("stack %s not supported", stack)
}
}
func (c *composeOperationsImpl) getMessages(stack ProductStack, operation ProcessorOperation) SubsystemOperationMessage {
msgs := SubsystemOperationMessages[SubsystemCompose][operation]
return SubsystemOperationMessage{
Enter: fmt.Sprintf(msgs.Enter, stack),
Exit: fmt.Sprintf(msgs.Exit, stack),
Error: fmt.Sprintf(msgs.Error, stack),
}
}
backend/cmd/installer/processor/docker.go
+432
View File
@@ -0,0 +1,432 @@
package processor
import (
"context"
"fmt"
"os"
"os/exec"
"path/filepath"
"strings"
"pentagi/pkg/tools"
cerrdefs "github.com/containerd/errdefs"
"github.com/docker/docker/api/types/container"
"github.com/docker/docker/api/types/image"
"github.com/docker/docker/api/types/network"
"github.com/docker/docker/api/types/volume"
"github.com/docker/docker/client"
)
type dockerOperationsImpl struct {
processor *processor
}
func newDockerOperations(p *processor) dockerOperations {
return &dockerOperationsImpl{processor: p}
}
func (d *dockerOperationsImpl) pullWorkerImage(ctx context.Context, state *operationState) error {
return d.pullImage(ctx, state, d.getWorkerImageName())
}
func (d *dockerOperationsImpl) pullDefaultImage(ctx context.Context, state *operationState) error {
return d.pullImage(ctx, state, d.getDefaultImageName())
}
func (d *dockerOperationsImpl) pullImage(ctx context.Context, state *operationState, name string) error {
d.processor.appendLog(fmt.Sprintf(MsgPullingImage, name), ProductStackWorker, state)
cmd := exec.CommandContext(ctx, "docker", "pull", name)
cmd.Env = d.getWorkerDockerEnv()
if err := d.processor.runCommand(cmd, ProductStackWorker, state); err != nil {
d.processor.appendLog(fmt.Sprintf(MsgImagePullFailed, name, err), ProductStackWorker, state)
return fmt.Errorf("failed to pull image %s: %w", name, err)
}
d.processor.appendLog(fmt.Sprintf(MsgImagePullCompleted, name), ProductStackWorker, state)
return nil
}
func (d *dockerOperationsImpl) removeWorkerContainers(ctx context.Context, state *operationState) error {
cli, err := d.createWorkerDockerClient()
if err != nil {
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
d.processor.appendLog(MsgRemovingWorkerContainers, ProductStackWorker, state)
allContainers, err := cli.ContainerList(ctx, container.ListOptions{All: true})
if err != nil {
return fmt.Errorf("failed to list worker containers: %w", err)
}
var containers []container.Summary
for _, c := range allContainers {
for _, name := range c.Names {
if strings.HasPrefix(name, "pentagi-") {
containers = append(containers, c)
break
}
}
}
if len(containers) == 0 {
d.processor.appendLog(MsgNoWorkerContainersFound, ProductStackWorker, state)
return nil
}
totalContainers := len(containers)
for _, cont := range containers {
if cont.State == "running" {
d.processor.appendLog(fmt.Sprintf(MsgStoppingContainer, cont.ID[:12]), ProductStackWorker, state)
if err := cli.ContainerStop(ctx, cont.ID, container.StopOptions{}); err != nil {
return fmt.Errorf("failed to stop worker container %s: %w", cont.ID, err)
}
}
d.processor.appendLog(fmt.Sprintf(MsgRemovingContainer, cont.ID[:12]), ProductStackWorker, state)
if err := cli.ContainerRemove(ctx, cont.ID, container.RemoveOptions{
Force: true,
}); err != nil {
return fmt.Errorf("failed to remove worker container %s: %w", cont.ID, err)
}
d.processor.appendLog(fmt.Sprintf(MsgContainerRemoved, cont.ID[:12]), ProductStackWorker, state)
}
d.processor.appendLog(fmt.Sprintf(MsgWorkerContainersRemoved, totalContainers), ProductStackWorker, state)
return nil
}
func (d *dockerOperationsImpl) removeWorkerImages(ctx context.Context, state *operationState) error {
return d.removeImages(ctx, state, image.RemoveOptions{
Force: false,
PruneChildren: false,
})
}
func (d *dockerOperationsImpl) purgeWorkerImages(ctx context.Context, state *operationState) error {
return d.removeImages(ctx, state, image.RemoveOptions{
Force: true,
PruneChildren: true,
})
}
func (d *dockerOperationsImpl) removeImages(
ctx context.Context, state *operationState, options image.RemoveOptions,
) error {
if err := d.removeWorkerContainers(ctx, state); err != nil {
return err
}
cli, err := d.createWorkerDockerClient()
if err != nil {
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
for _, imageName := range []string{d.getWorkerImageName(), d.getDefaultImageName()} {
d.processor.appendLog(fmt.Sprintf(MsgRemovingImage, imageName), ProductStackWorker, state)
if _, err := cli.ImageRemove(ctx, imageName, options); err != nil {
if !cerrdefs.IsNotFound(err) {
return fmt.Errorf("failed to remove image %s: %w", imageName, err)
} else {
d.processor.appendLog(fmt.Sprintf(MsgImageNotFound, imageName), ProductStackWorker, state)
}
} else {
d.processor.appendLog(fmt.Sprintf(MsgImageRemoved, imageName), ProductStackWorker, state)
}
}
d.processor.appendLog(MsgWorkerImagesRemoveCompleted, ProductStackWorker, state)
return nil
}
// createMainDockerClient creates docker client for the main stack (non-worker) using current process env
func (d *dockerOperationsImpl) createMainDockerClient() (*client.Client, error) {
return client.NewClientWithOpts(
client.FromEnv,
client.WithAPIVersionNegotiation(),
)
}
// checkMainDockerNetwork returns true if a docker network with given name exists
func (d *dockerOperationsImpl) checkMainDockerNetwork(ctx context.Context, cli *client.Client, name string) (bool, error) {
nets, err := cli.NetworkList(ctx, network.ListOptions{})
if err != nil {
return false, fmt.Errorf("failed to list docker networks: %w", err)
}
for _, n := range nets {
if n.Name == name {
return true, nil
}
}
return false, nil
}
// createMainDockerNetwork creates a docker network with given name if it does not exist
func (d *dockerOperationsImpl) createMainDockerNetwork(ctx context.Context, cli *client.Client, state *operationState, name string) error {
exists, err := d.checkMainDockerNetwork(ctx, cli, name)
if err != nil {
return err
}
if exists {
// inspect to validate labels
nw, err := cli.NetworkInspect(ctx, name, network.InspectOptions{})
if err == nil {
wantProject := ""
if envPath := d.processor.state.GetEnvPath(); envPath != "" {
wantProject = filepath.Base(filepath.Dir(envPath))
}
hasComposeLabel := nw.Labels["com.docker.compose.network"] == name
projectMatches := wantProject == "" || nw.Labels["com.docker.compose.project"] == wantProject
if hasComposeLabel && projectMatches {
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkExists, name), ProductStackInstaller, state)
return nil
}
// if labels incorrect and network has no containers attached, recreate with correct labels
if len(nw.Containers) > 0 {
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkInUse, name), ProductStackInstaller, state)
return nil
}
d.processor.appendLog(fmt.Sprintf(MsgRecreatingDockerNetwork, name), ProductStackInstaller, state)
if err := cli.NetworkRemove(ctx, nw.ID); err != nil {
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkRemoveFailed, name, err), ProductStackInstaller, state)
return fmt.Errorf("failed to remove network %s: %w", name, err)
}
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkRemoved, name), ProductStackInstaller, state)
}
}
d.processor.appendLog(fmt.Sprintf(MsgCreatingDockerNetwork, name), ProductStackInstaller, state)
// mimic docker compose-created network by setting compose labels
// project name: derived from working directory of env file (same as compose default)
projectName := ""
if envPath := d.processor.state.GetEnvPath(); envPath != "" {
projectName = filepath.Base(filepath.Dir(envPath))
}
labels := map[string]string{
"com.docker.compose.network": name,
}
if projectName != "" {
labels["com.docker.compose.project"] = projectName
}
// driver: bridge (compose default for local networks)
_, err = cli.NetworkCreate(ctx, name, network.CreateOptions{
Driver: "bridge",
Labels: labels,
})
if err != nil {
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkCreateFailed, name, err), ProductStackInstaller, state)
return fmt.Errorf("failed to create network %s: %w", name, err)
}
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkCreated, name), ProductStackInstaller, state)
return nil
}
// ensureMainDockerNetworks ensures all required networks for main stacks exist
func (d *dockerOperationsImpl) ensureMainDockerNetworks(ctx context.Context, state *operationState) error {
d.processor.appendLog(MsgEnsuringDockerNetworks, ProductStackInstaller, state)
defer d.processor.appendLog("", ProductStackInstaller, state)
if !d.processor.checker.DockerApiAccessible {
return fmt.Errorf("docker api is not accessible")
}
cli, err := d.createMainDockerClient()
if err != nil {
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
required := []string{
string(ProductDockerNetworkPentagi),
string(ProductDockerNetworkObservability),
string(ProductDockerNetworkLangfuse),
}
for _, net := range required {
if err := d.createMainDockerNetwork(ctx, cli, state, net); err != nil {
return err
}
}
return nil
}
// removeMainDockerNetwork removes a docker network by name, detaching containers if possible
func (d *dockerOperationsImpl) removeMainDockerNetwork(ctx context.Context, state *operationState, name string) error {
if !d.processor.checker.DockerApiAccessible {
return fmt.Errorf("docker api is not accessible")
}
cli, err := d.createMainDockerClient()
if err != nil {
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
// try inspect; if not found just return
nw, err := cli.NetworkInspect(ctx, name, network.InspectOptions{})
if err != nil {
return nil
}
// attempt to disconnect all containers first (best effort)
for id := range nw.Containers {
_ = cli.NetworkDisconnect(ctx, nw.ID, id, true)
}
if err := cli.NetworkRemove(ctx, nw.ID); err != nil {
return err
}
d.processor.appendLog(fmt.Sprintf(MsgDockerNetworkRemoved, name), ProductStackInstaller, state)
return nil
}
// removeMainImages removes a list of images from main docker daemon
func (d *dockerOperationsImpl) removeMainImages(ctx context.Context, state *operationState, images []string) error {
if !d.processor.checker.DockerApiAccessible {
return fmt.Errorf("docker api is not accessible")
}
cli, err := d.createMainDockerClient()
if err != nil {
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
opts := image.RemoveOptions{Force: state.force, PruneChildren: state.force}
for _, img := range images {
if img == "" {
continue
}
d.processor.appendLog(fmt.Sprintf(MsgRemovingImage, img), ProductStackInstaller, state)
if _, err := cli.ImageRemove(ctx, img, opts); err != nil {
if !cerrdefs.IsNotFound(err) {
return err
}
}
}
return nil
}
// removeWorkerVolumes removes worker volumes (pentagi-terminal-*-data) in worker environment
func (d *dockerOperationsImpl) removeWorkerVolumes(ctx context.Context, state *operationState) error {
cli, err := d.createWorkerDockerClient()
if err != nil {
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
vols, err := cli.VolumeList(ctx, volume.ListOptions{})
if err != nil {
return err
}
for _, v := range vols.Volumes {
if strings.HasPrefix(v.Name, tools.PrimaryTerminalNamePrefix) && strings.HasSuffix(v.Name, "-data") {
_ = cli.VolumeRemove(ctx, v.Name, true)
}
}
return nil
}
func (d *dockerOperationsImpl) createWorkerDockerClient() (*client.Client, error) {
opts := []client.Opt{
client.WithAPIVersionNegotiation(),
}
envVar, exists := d.processor.state.GetVar(client.EnvOverrideHost)
if exists && (envVar.Value != "" || envVar.IsChanged) {
opts = append(opts, client.WithHost(envVar.Value))
} else if envVar.Default != "" {
opts = append(opts, client.WithHost(envVar.Default))
} else if envVar := os.Getenv(client.EnvOverrideHost); envVar != "" {
opts = append(opts, client.WithHost(envVar))
} else {
opts = append(opts, client.WithHostFromEnv())
}
type tlsConfig struct {
certPath string
keyPath string
caPath string
}
getTLSConfig := func(path string) tlsConfig {
return tlsConfig{
certPath: filepath.Join(path, "cert.pem"),
keyPath: filepath.Join(path, "key.pem"),
caPath: filepath.Join(path, "ca.pem"),
}
}
envVar, exists = d.processor.state.GetVar("PENTAGI_" + client.EnvOverrideCertPath)
if exists && (envVar.Value != "" || envVar.IsChanged) {
cfg := getTLSConfig(envVar.Value)
opts = append(opts, client.WithTLSClientConfig(cfg.certPath, cfg.keyPath, cfg.caPath))
} else if envVar.Default != "" {
cfg := getTLSConfig(envVar.Default)
opts = append(opts, client.WithTLSClientConfig(cfg.certPath, cfg.keyPath, cfg.caPath))
} else {
opts = append(opts, client.WithTLSClientConfigFromEnv())
}
return client.NewClientWithOpts(opts...)
}
func (d *dockerOperationsImpl) getWorkerDockerEnv() []string {
var env []string
envVar, exists := d.processor.state.GetVar(client.EnvOverrideHost)
if exists && (envVar.Value != "" || envVar.IsChanged) {
env = append(env, fmt.Sprintf("%s=%s", client.EnvOverrideHost, envVar.Value))
} else if envVar.Default != "" {
env = append(env, fmt.Sprintf("%s=%s", client.EnvOverrideHost, envVar.Default))
} else if envVar := os.Getenv(client.EnvOverrideHost); envVar != "" {
env = append(env, fmt.Sprintf("%s=%s", client.EnvOverrideHost, envVar))
}
envVar, exists = d.processor.state.GetVar("PENTAGI_" + client.EnvOverrideCertPath)
if exists && (envVar.Value != "" || envVar.IsChanged) {
env = append(env, fmt.Sprintf("%s=%s", client.EnvOverrideCertPath, envVar.Value))
} else if envVar.Default != "" {
env = append(env, fmt.Sprintf("%s=%s", client.EnvOverrideCertPath, envVar.Default))
} else if envVar := os.Getenv(client.EnvOverrideCertPath); envVar != "" {
env = append(env, fmt.Sprintf("%s=%s", client.EnvOverrideCertPath, envVar))
}
envVar, exists = d.processor.state.GetVar(client.EnvTLSVerify)
if exists && (envVar.Value != "" || envVar.IsChanged) {
env = append(env, fmt.Sprintf("%s=%s", client.EnvTLSVerify, envVar.Value))
} else if envVar.Default != "" {
env = append(env, fmt.Sprintf("%s=%s", client.EnvTLSVerify, envVar.Default))
} else if envVar := os.Getenv(client.EnvTLSVerify); envVar != "" {
env = append(env, fmt.Sprintf("%s=%s", client.EnvTLSVerify, envVar))
}
return env
}
func (d *dockerOperationsImpl) getWorkerImageName() string {
envVar, exists := d.processor.state.GetVar("DOCKER_DEFAULT_IMAGE_FOR_PENTEST")
if exists && envVar.Value != "" {
return envVar.Value
}
if envVar.Default != "" {
return envVar.Default
}
return "vxcontrol/kali-linux:latest"
}
func (d *dockerOperationsImpl) getDefaultImageName() string {
envVar, exists := d.processor.state.GetVar("DOCKER_DEFAULT_IMAGE")
if exists && envVar.Value != "" {
return envVar.Value
}
if envVar.Default != "" {
return envVar.Default
}
return "debian:latest"
}
backend/cmd/installer/processor/fs.go
+427
View File
@@ -0,0 +1,427 @@
package processor
import (
"context"
"errors"
"fmt"
"maps"
"os"
"path/filepath"
"pentagi/cmd/installer/files"
"gopkg.in/yaml.v3"
)
const (
observabilityDirectory = "observability"
pentagiExampleCustomConfigLLM = "example.custom.provider.yml"
pentagiExampleOllamaConfigLLM = "example.ollama.provider.yml"
)
var filesToExcludeFromVerification = []string{
"observability/otel/config.yml",
"observability/grafana/config/grafana.ini",
pentagiExampleCustomConfigLLM,
pentagiExampleOllamaConfigLLM,
}
var allStacks = []ProductStack{
ProductStackPentagi,
ProductStackGraphiti,
ProductStackLangfuse,
ProductStackObservability,
}
type fileSystemOperationsImpl struct {
processor *processor
}
func newFileSystemOperations(p *processor) fileSystemOperations {
return &fileSystemOperationsImpl{processor: p}
}
func (fs *fileSystemOperationsImpl) ensureStackIntegrity(ctx context.Context, stack ProductStack, state *operationState) error {
fs.processor.appendLog(fmt.Sprintf(MsgEnsurngStackIntegrity, stack), stack, state)
defer fs.processor.appendLog("", stack, state)
switch stack {
case ProductStackPentagi:
errCompose := fs.ensureFileFromEmbed(composeFilePentagi, state)
errCustom := fs.ensureFileFromEmbed(pentagiExampleCustomConfigLLM, state)
errOllama := fs.ensureFileFromEmbed(pentagiExampleOllamaConfigLLM, state)
return errors.Join(errCompose, errCustom, errOllama)
case ProductStackGraphiti:
return fs.ensureFileFromEmbed(composeFileGraphiti, state)
case ProductStackLangfuse:
return fs.ensureFileFromEmbed(composeFileLangfuse, state)
case ProductStackObservability:
errCompose := fs.ensureFileFromEmbed(composeFileObservability, state)
errDirectory := fs.ensureDirectoryFromEmbed(observabilityDirectory, state)
return errors.Join(errCompose, errDirectory)
case ProductStackAll, ProductStackCompose:
// process all stacks sequentially
for _, s := range allStacks {
if err := fs.ensureStackIntegrity(ctx, s, state); err != nil {
return err
}
}
return nil
default:
return fmt.Errorf("operation ensure integrity not applicable for stack %s", stack)
}
}
func (fs *fileSystemOperationsImpl) verifyStackIntegrity(ctx context.Context, stack ProductStack, state *operationState) error {
fs.processor.appendLog(fmt.Sprintf(MsgVerifyingStackIntegrity, stack), stack, state)
defer fs.processor.appendLog("", stack, state)
switch stack {
case ProductStackPentagi:
return fs.verifyFileIntegrity(composeFilePentagi, state)
case ProductStackGraphiti:
return fs.verifyFileIntegrity(composeFileGraphiti, state)
case ProductStackLangfuse:
return fs.verifyFileIntegrity(composeFileLangfuse, state)
case ProductStackObservability:
if err := fs.verifyFileIntegrity(composeFileObservability, state); err != nil {
return err
}
return fs.verifyDirectoryIntegrity(observabilityDirectory, state)
case ProductStackAll, ProductStackCompose:
// process all stacks sequentially
for _, s := range allStacks {
if err := fs.verifyStackIntegrity(ctx, s, state); err != nil {
return err
}
}
return nil
default:
return fmt.Errorf("operation verify integrity not applicable for stack %s", stack)
}
}
// checkStackIntegrity is a silent version of verifyStackIntegrity, used for getting files statuses
func (fs *fileSystemOperationsImpl) checkStackIntegrity(ctx context.Context, stack ProductStack) (FilesCheckResult, error) {
result := make(FilesCheckResult)
switch stack {
case ProductStackPentagi:
result[composeFilePentagi] = fs.checkFileIntegrity(composeFilePentagi)
case ProductStackGraphiti:
result[composeFileGraphiti] = fs.checkFileIntegrity(composeFileGraphiti)
case ProductStackLangfuse:
result[composeFileLangfuse] = fs.checkFileIntegrity(composeFileLangfuse)
case ProductStackObservability:
result[composeFileObservability] = fs.checkFileIntegrity(composeFileObservability)
if r, err := fs.checkDirectoryIntegrity(observabilityDirectory); err != nil {
return result, err
} else {
maps.Copy(result, r)
}
case ProductStackAll, ProductStackCompose:
// process all stacks sequentially
for _, s := range allStacks {
if r, err := fs.checkStackIntegrity(ctx, s); err != nil {
return result, err // early exit after first error
} else {
maps.Copy(result, r)
}
}
default:
return result, fmt.Errorf("operation check integrity not applicable for stack %s", stack)
}
return result, nil
}
func (fs *fileSystemOperationsImpl) cleanupStackFiles(ctx context.Context, stack ProductStack, state *operationState) error {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
fs.processor.appendLog(fmt.Sprintf(MsgCleaningUpStackFiles, stack), stack, state)
defer fs.processor.appendLog("", stack, state)
var filesToRemove []string
switch stack {
case ProductStackPentagi:
filesToRemove = append(filesToRemove, filepath.Join(workingDir, composeFilePentagi))
filesToRemove = append(filesToRemove, filepath.Join(workingDir, pentagiExampleCustomConfigLLM))
filesToRemove = append(filesToRemove, filepath.Join(workingDir, pentagiExampleOllamaConfigLLM))
case ProductStackGraphiti:
filesToRemove = append(filesToRemove, filepath.Join(workingDir, composeFileGraphiti))
case ProductStackLangfuse:
filesToRemove = append(filesToRemove, filepath.Join(workingDir, composeFileLangfuse))
case ProductStackObservability:
filesToRemove = append(filesToRemove, filepath.Join(workingDir, composeFileObservability))
filesToRemove = append(filesToRemove, filepath.Join(workingDir, observabilityDirectory))
case ProductStackAll, ProductStackCompose:
for _, s := range allStacks {
if err := fs.cleanupStackFiles(ctx, s, state); err != nil {
return err
}
}
return nil
default:
return fmt.Errorf("operation cleanup not applicable for stack %s", stack)
}
for _, path := range filesToRemove {
if err := os.RemoveAll(path); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove %s: %w", path, err)
}
}
fs.processor.appendLog(MsgStackFilesCleanupCompleted, stack, state)
return nil
}
func (fs *fileSystemOperationsImpl) ensureFileFromEmbed(filename string, state *operationState) error {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
targetPath := filepath.Join(workingDir, filename)
if !state.force && fs.fileExists(targetPath) {
fs.processor.appendLog(fmt.Sprintf(MsgFileIntegrityValid, filename), ProductStackAll, state)
return nil
}
if fs.fileExists(targetPath) {
fs.processor.appendLog(fmt.Sprintf(MsgUpdatingExistingFile, filename), ProductStackAll, state)
} else {
fs.processor.appendLog(fmt.Sprintf(MsgCreatingMissingFile, filename), ProductStackAll, state)
}
return fs.processor.files.Copy(filename, workingDir, true)
}
func (fs *fileSystemOperationsImpl) ensureDirectoryFromEmbed(dirname string, state *operationState) error {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
targetPath := filepath.Join(workingDir, dirname)
if !state.force && fs.directoryExists(targetPath) {
return fs.verifyDirectoryContentIntegrity(dirname, targetPath, state)
}
if fs.directoryExists(targetPath) {
fs.processor.appendLog(fmt.Sprintf(MsgUpdatingExistingFile, dirname), ProductStackAll, state)
} else {
fs.processor.appendLog(fmt.Sprintf(MsgCreatingMissingFile, dirname), ProductStackAll, state)
}
return fs.processor.files.Copy(dirname, workingDir, true)
}
func (fs *fileSystemOperationsImpl) checkFileIntegrity(filename string) files.FileStatus {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
return fs.processor.files.Check(filename, workingDir)
}
func (fs *fileSystemOperationsImpl) checkDirectoryIntegrity(dirname string) (FilesCheckResult, error) {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
targetPath := filepath.Join(workingDir, dirname)
return fs.checkDirectoryContentIntegrity(dirname, targetPath)
}
func (fs *fileSystemOperationsImpl) checkDirectoryContentIntegrity(embedPath, targetPath string) (FilesCheckResult, error) {
if !fs.processor.files.Exists(embedPath) {
return FilesCheckResult{embedPath: files.FileStatusMissing}, nil
}
info, err := fs.processor.files.Stat(embedPath)
if err != nil {
return FilesCheckResult{}, fmt.Errorf("failed to stat embedded directory %s: %w", embedPath, err)
}
if !info.IsDir() {
return FilesCheckResult{embedPath: fs.checkFileIntegrity(embedPath)}, nil
}
// get list of embedded files in directory
embeddedFiles, err := fs.processor.files.List(embedPath)
if err != nil {
return FilesCheckResult{}, fmt.Errorf("failed to list embedded files in %s: %w", embedPath, err)
}
// check each embedded file exists and matches in target directory except excluded files
result := make(FilesCheckResult)
workingDir := filepath.Dir(targetPath)
for _, embeddedFile := range embeddedFiles {
status := fs.processor.files.Check(embeddedFile, workingDir)
// skip integrity tracking for excluded files but ensure their presence
if !fs.isExcludedFromVerification(embeddedFile) || status != files.FileStatusModified {
result[embeddedFile] = status
}
}
return result, nil
}
func (fs *fileSystemOperationsImpl) verifyFileIntegrity(filename string, state *operationState) error {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
targetPath := filepath.Join(workingDir, filename)
if !fs.fileExists(targetPath) {
fs.processor.appendLog(fmt.Sprintf(MsgCreatingMissingFile, filename), ProductStackAll, state)
return fs.processor.files.Copy(filename, workingDir, true)
}
if state.force {
fs.processor.appendLog(fmt.Sprintf(MsgUpdatingExistingFile, filename), ProductStackAll, state)
return fs.processor.files.Copy(filename, workingDir, true)
}
if err := fs.validateYamlFile(targetPath); err != nil {
fs.processor.appendLog(fmt.Sprintf(MsgUpdatingExistingFile, filename), ProductStackAll, state)
return fs.processor.files.Copy(filename, workingDir, true)
}
fs.processor.appendLog(fmt.Sprintf(MsgFileIntegrityValid, filename), ProductStackAll, state)
return nil
}
func (fs *fileSystemOperationsImpl) verifyDirectoryIntegrity(dirname string, state *operationState) error {
workingDir := filepath.Dir(fs.processor.state.GetEnvPath())
targetPath := filepath.Join(workingDir, dirname)
if !fs.directoryExists(targetPath) {
fs.processor.appendLog(fmt.Sprintf(MsgCreatingMissingFile, dirname), ProductStackAll, state)
return fs.processor.files.Copy(dirname, workingDir, true)
}
if state.force {
// update directory content selectively, respecting excluded files
fs.processor.appendLog(fmt.Sprintf(MsgUpdatingExistingFile, dirname), ProductStackAll, state)
return fs.verifyDirectoryContentIntegrity(dirname, targetPath, state)
}
return fs.verifyDirectoryContentIntegrity(dirname, targetPath, state)
}
func (fs *fileSystemOperationsImpl) verifyDirectoryContentIntegrity(embedPath, targetPath string, state *operationState) error {
if !fs.processor.files.Exists(embedPath) {
return fmt.Errorf("embedded directory %s not found", embedPath)
}
info, err := fs.processor.files.Stat(embedPath)
if err != nil {
return fmt.Errorf("failed to stat embedded directory %s: %w", embedPath, err)
}
if !info.IsDir() {
return fs.verifyFileIntegrity(embedPath, state)
}
// get list of embedded files in directory
embeddedFiles, err := fs.processor.files.List(embedPath)
if err != nil {
return fmt.Errorf("failed to list embedded files in %s: %w", embedPath, err)
}
// verify each embedded file exists and matches in target directory
// note: targetPath is the full path to the directory we're verifying
// we need to check files relative to the parent of targetPath
workingDir := filepath.Dir(targetPath)
hasUnupdatedModified := false
for _, embeddedFile := range embeddedFiles {
// skip integrity tracking for excluded files but ensure their presence
if fs.isExcludedFromVerification(embeddedFile) {
status := fs.processor.files.Check(embeddedFile, workingDir)
if status == files.FileStatusMissing {
fs.processor.appendLog(fmt.Sprintf(MsgCreatingMissingFile, embeddedFile), ProductStackAll, state)
if err := fs.processor.files.Copy(embeddedFile, workingDir, true); err != nil {
return fmt.Errorf("failed to copy missing excluded file %s: %w", embeddedFile, err)
}
}
// do not mark as modified and do not overwrite if present
continue
}
// check file status using optimized hash comparison
status := fs.processor.files.Check(embeddedFile, workingDir)
switch status {
case files.FileStatusMissing:
fs.processor.appendLog(fmt.Sprintf(MsgCreatingMissingFile, embeddedFile), ProductStackAll, state)
if err := fs.processor.files.Copy(embeddedFile, workingDir, true); err != nil {
return fmt.Errorf("failed to copy missing file %s: %w", embeddedFile, err)
}
case files.FileStatusModified:
if state.force {
fs.processor.appendLog(fmt.Sprintf(MsgUpdatingExistingFile, embeddedFile), ProductStackAll, state)
if err := fs.processor.files.Copy(embeddedFile, workingDir, true); err != nil {
return fmt.Errorf("failed to update modified file %s: %w", embeddedFile, err)
}
} else {
hasUnupdatedModified = true
fs.processor.appendLog(fmt.Sprintf(MsgSkippingModifiedFile, embeddedFile), ProductStackAll, state)
}
case files.FileStatusOK:
// file is valid, no action needed
}
}
if hasUnupdatedModified {
fs.processor.appendLog(fmt.Sprintf(MsgDirectoryCheckedWithModified, embedPath), ProductStackAll, state)
} else {
fs.processor.appendLog(fmt.Sprintf(MsgFileIntegrityValid, embedPath), ProductStackAll, state)
}
return nil
}
// isExcludedFromVerification returns true if the provided path should be excluded
// from integrity verification. The file should still exist on disk, but its
// content modifications must not trigger updates or verification failures.
func (fs *fileSystemOperationsImpl) isExcludedFromVerification(path string) bool {
normalized := filepath.ToSlash(path)
for _, excluded := range filesToExcludeFromVerification {
if filepath.ToSlash(excluded) == normalized {
return true
}
}
return false
}
func (fs *fileSystemOperationsImpl) fileExists(path string) bool {
info, err := os.Stat(path)
return err == nil && !info.IsDir()
}
func (fs *fileSystemOperationsImpl) directoryExists(path string) bool {
info, err := os.Stat(path)
return err == nil && info.IsDir()
}
func (fs *fileSystemOperationsImpl) validateYamlFile(path string) error {
content, err := os.ReadFile(path)
if err != nil {
return fmt.Errorf("failed to read file: %w", err)
}
var data interface{}
if err := yaml.Unmarshal(content, &data); err != nil {
return fmt.Errorf("invalid YAML syntax: %w", err)
}
return nil
}
backend/cmd/installer/processor/fs_test.go
+712
View File
@@ -0,0 +1,712 @@
package processor
import (
"context"
"os"
"path/filepath"
"sync"
"testing"
"pentagi/cmd/installer/files"
)
// testStackIntegrityOperation is a helper for testing stack integrity operations
func testStackIntegrityOperation(t *testing.T, operation func(*fileSystemOperationsImpl, context.Context, ProductStack, *operationState) error, needsTempDir bool) {
t.Helper()
tests := []struct {
name string
stack ProductStack
expectErr bool
}{
{"ProductStackPentagi", ProductStackPentagi, false},
{"ProductStackLangfuse", ProductStackLangfuse, false},
{"ProductStackObservability", ProductStackObservability, false},
{"ProductStackCompose", ProductStackCompose, false},
{"ProductStackAll", ProductStackAll, false},
{"ProductStackWorker - unsupported", ProductStackWorker, true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
processor := createTestProcessor()
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
// for cleanup operations, ensure temp directory setup
if needsTempDir {
tmpDir := t.TempDir()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
}
err := operation(fsOps, t.Context(), tt.stack, testOperationState(t))
assertError(t, err, tt.expectErr, "")
})
}
}
func TestFileSystemOperationsImpl_EnsureStackIntegrity(t *testing.T) {
testStackIntegrityOperation(t, (*fileSystemOperationsImpl).ensureStackIntegrity, false)
}
func TestFileSystemOperationsImpl_VerifyStackIntegrity(t *testing.T) {
testStackIntegrityOperation(t, (*fileSystemOperationsImpl).verifyStackIntegrity, false)
}
func TestFileSystemOperationsImpl_CleanupStackFiles(t *testing.T) {
testStackIntegrityOperation(t, (*fileSystemOperationsImpl).cleanupStackFiles, true)
}
func TestFileSystemOperationsImpl_EnsureFileFromEmbed(t *testing.T) {
tests := []struct {
name string
filename string
force bool
setup func(*mockFiles, string) // setup mock and working dir
}{
{
name: "file missing - should copy",
filename: "test.yml",
force: false,
setup: func(m *mockFiles, workingDir string) {
// file not exists, will be copied
m.AddFile("test.yml", []byte("test content"))
},
},
{
name: "file exists, force false - should skip",
filename: "test.yml",
force: false,
setup: func(m *mockFiles, workingDir string) {
// create existing file
os.WriteFile(filepath.Join(workingDir, "test.yml"), []byte("existing"), 0644)
m.AddFile("test.yml", []byte("test content"))
},
},
{
name: "file exists, force true - should update",
filename: "test.yml",
force: true,
setup: func(m *mockFiles, workingDir string) {
// create existing file
os.WriteFile(filepath.Join(workingDir, "test.yml"), []byte("existing"), 0644)
m.AddFile("test.yml", []byte("test content"))
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create temp directory
tmpDir := t.TempDir()
// create test processor using unified mocks
processor := createTestProcessor()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
mockFiles := processor.files.(*mockFiles)
tt.setup(mockFiles, tmpDir)
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
state := &operationState{force: tt.force, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.ensureFileFromEmbed(tt.filename, state)
assertNoError(t, err)
})
}
}
func TestFileSystemOperationsImpl_VerifyDirectoryContentIntegrity(t *testing.T) {
tests := []struct {
name string
embedPath string
setup func(*mockFiles, string, string) // setup mock, embedPath, targetPath
force bool
expectErr bool
}{
{
name: "embedded directory not found",
embedPath: "nonexistent",
setup: func(m *mockFiles, embedPath, targetPath string) {},
force: false,
expectErr: true,
},
{
name: "directory with all files OK",
embedPath: "observability",
setup: func(m *mockFiles, embedPath, targetPath string) {
// setup embedded files
m.lists[embedPath] = []string{
"observability/config1.yml",
"observability/config2.yml",
}
m.statuses["observability/config1.yml"] = files.FileStatusOK
m.statuses["observability/config2.yml"] = files.FileStatusOK
m.AddFile(embedPath, []byte{}) // mark as existing directory
// create target directory
os.MkdirAll(targetPath, 0755)
},
force: false,
expectErr: false,
},
{
name: "directory with missing files",
embedPath: "observability",
setup: func(m *mockFiles, embedPath, targetPath string) {
// setup embedded files
m.lists[embedPath] = []string{
"observability/config1.yml",
"observability/config2.yml",
}
m.statuses["observability/config1.yml"] = files.FileStatusOK
m.statuses["observability/config2.yml"] = files.FileStatusMissing
m.AddFile(embedPath, []byte{}) // mark as existing directory
// create target directory
os.MkdirAll(targetPath, 0755)
},
force: false,
expectErr: false,
},
{
name: "directory with modified files, force false",
embedPath: "observability",
setup: func(m *mockFiles, embedPath, targetPath string) {
// setup embedded files
m.lists[embedPath] = []string{
"observability/config1.yml",
}
m.statuses["observability/config1.yml"] = files.FileStatusModified
m.AddFile(embedPath, []byte{}) // mark as existing directory
// create target directory
os.MkdirAll(targetPath, 0755)
},
force: false,
expectErr: false,
},
{
name: "directory with modified files, force true",
embedPath: "observability",
setup: func(m *mockFiles, embedPath, targetPath string) {
// setup embedded files
m.lists[embedPath] = []string{
"observability/config1.yml",
}
m.statuses["observability/config1.yml"] = files.FileStatusModified
m.AddFile(embedPath, []byte{}) // mark as existing directory
// create target directory
os.MkdirAll(targetPath, 0755)
},
force: true,
expectErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create temp directory
tmpDir := t.TempDir()
targetPath := filepath.Join(tmpDir, tt.embedPath)
// create test processor
processor := createTestProcessor()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
mockFiles := processor.files.(*mockFiles)
tt.setup(mockFiles, tt.embedPath, targetPath)
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
state := &operationState{force: tt.force, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.verifyDirectoryContentIntegrity(tt.embedPath, targetPath, state)
assertError(t, err, tt.expectErr, "")
})
}
}
func TestFileSystemOperationsImpl_ExcludedFilesHandling(t *testing.T) {
if len(filesToExcludeFromVerification) == 0 {
t.Skip("no excluded files configured; skipping excluded files tests")
}
excluded := filesToExcludeFromVerification[0]
t.Run("excluded_missing_should_be_copied", func(t *testing.T) {
tmpDir := t.TempDir()
targetPath := filepath.Join(tmpDir, observabilityDirectory)
processor := createTestProcessor()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
mockFiles := processor.files.(*mockFiles)
// mark embedded directory and list with excluded file only
mockFiles.lists[observabilityDirectory] = []string{
excluded,
}
mockFiles.statuses[excluded] = files.FileStatusMissing
mockFiles.AddFile(observabilityDirectory, []byte{})
// ensure target directory exists on fs
_ = os.MkdirAll(targetPath, 0o755)
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
state := &operationState{force: false, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.verifyDirectoryContentIntegrity(observabilityDirectory, targetPath, state)
assertNoError(t, err)
if len(mockFiles.copies) != 1 {
t.Fatalf("expected 1 copy for missing excluded file, got %d", len(mockFiles.copies))
}
if mockFiles.copies[0].Src != excluded || mockFiles.copies[0].Dst != tmpDir {
t.Errorf("unexpected copy details: %+v", mockFiles.copies[0])
}
})
t.Run("excluded_modified_should_not_be_copied", func(t *testing.T) {
tmpDir := t.TempDir()
targetPath := filepath.Join(tmpDir, observabilityDirectory)
processor := createTestProcessor()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
mockFiles := processor.files.(*mockFiles)
// mark embedded directory and list with excluded file only
mockFiles.lists[observabilityDirectory] = []string{
excluded,
}
mockFiles.statuses[excluded] = files.FileStatusModified
mockFiles.AddFile(observabilityDirectory, []byte{})
// ensure target directory exists on fs
_ = os.MkdirAll(targetPath, 0o755)
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
state := &operationState{force: false, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.verifyDirectoryContentIntegrity(observabilityDirectory, targetPath, state)
assertNoError(t, err)
if len(mockFiles.copies) != 0 {
t.Fatalf("expected 0 copies for modified excluded file, got %d", len(mockFiles.copies))
}
})
t.Run("force_true_updates_only_non_excluded", func(t *testing.T) {
tmpDir := t.TempDir()
targetPath := filepath.Join(tmpDir, observabilityDirectory)
processor := createTestProcessor()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
mockFiles := processor.files.(*mockFiles)
// list contains excluded and non-excluded files
nonExcluded := "observability/other.yml"
mockFiles.lists[observabilityDirectory] = []string{
excluded, // excluded
nonExcluded, // non-excluded
}
mockFiles.statuses[excluded] = files.FileStatusModified
mockFiles.statuses[nonExcluded] = files.FileStatusModified
mockFiles.AddFile(observabilityDirectory, []byte{})
// ensure target directory exists on fs
_ = os.MkdirAll(targetPath, 0o755)
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
// call higher-level method to exercise force=true branch
state := &operationState{force: true, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.verifyDirectoryIntegrity(observabilityDirectory, state)
assertNoError(t, err)
if len(mockFiles.copies) != 1 {
t.Fatalf("expected 1 copy for non-excluded modified file, got %d", len(mockFiles.copies))
}
if mockFiles.copies[0].Src != nonExcluded || mockFiles.copies[0].Dst != tmpDir {
t.Errorf("unexpected copy details: %+v", mockFiles.copies[0])
}
})
}
func TestFileSystemOperationsImpl_FileExists(t *testing.T) {
tmpDir := t.TempDir()
processor := createTestProcessor()
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
// create test file
testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("test"), 0644)
// create test directory
testDir := filepath.Join(tmpDir, "testdir")
os.MkdirAll(testDir, 0755)
tests := []struct {
name string
path string
expected bool
}{
{"existing file", testFile, true},
{"existing directory", testDir, false}, // fileExists should return false for directories
{"nonexistent path", filepath.Join(tmpDir, "nonexistent"), false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := fsOps.fileExists(tt.path)
if result != tt.expected {
t.Errorf("fileExists(%s) = %v, want %v", tt.path, result, tt.expected)
}
})
}
}
func TestFileSystemOperationsImpl_DirectoryExists(t *testing.T) {
tmpDir := t.TempDir()
processor := createTestProcessor()
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
// create test file
testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("test"), 0644)
// create test directory
testDir := filepath.Join(tmpDir, "testdir")
os.MkdirAll(testDir, 0755)
tests := []struct {
name string
path string
expected bool
}{
{"existing file", testFile, false}, // directoryExists should return false for files
{"existing directory", testDir, true},
{"nonexistent path", filepath.Join(tmpDir, "nonexistent"), false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := fsOps.directoryExists(tt.path)
if result != tt.expected {
t.Errorf("directoryExists(%s) = %v, want %v", tt.path, result, tt.expected)
}
})
}
}
func TestFileSystemOperationsImpl_ValidateYamlFile(t *testing.T) {
tmpDir := t.TempDir()
processor := createTestProcessor()
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
tests := []struct {
name string
content string
expectErr bool
}{
{
name: "valid YAML",
content: `
version: '3.8'
services:
app:
image: nginx
ports:
- "80:80"
`,
expectErr: false,
},
{
name: "invalid YAML - syntax error",
content: `
version: '3.8'
services:
app:
image: nginx
ports:
- "80:80
# missing closing quote
`,
expectErr: true,
},
{
name: "empty file",
content: "",
expectErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// create test file
testFile := filepath.Join(tmpDir, "test.yml")
err := os.WriteFile(testFile, []byte(tt.content), 0644)
if err != nil {
t.Fatalf("failed to create test file: %v", err)
}
err = fsOps.validateYamlFile(testFile)
assertError(t, err, tt.expectErr, "")
})
}
}
func TestCheckStackIntegrity(t *testing.T) {
tests := []struct {
name string
stack ProductStack
setup func(*mockFiles)
expected map[string]files.FileStatus
}{
{
name: "pentagi_stack_all_files_ok",
stack: ProductStackPentagi,
setup: func(m *mockFiles) {
m.statuses[composeFilePentagi] = files.FileStatusOK
},
expected: map[string]files.FileStatus{
composeFilePentagi: files.FileStatusOK,
},
},
{
name: "langfuse_stack_file_modified",
stack: ProductStackLangfuse,
setup: func(m *mockFiles) {
m.statuses[composeFileLangfuse] = files.FileStatusModified
},
expected: map[string]files.FileStatus{
composeFileLangfuse: files.FileStatusModified,
},
},
{
name: "observability_stack_mixed_status",
stack: ProductStackObservability,
setup: func(m *mockFiles) {
m.statuses[composeFileObservability] = files.FileStatusOK
m.lists[observabilityDirectory] = []string{
"observability/config1.yml",
"observability/config2.yml",
"observability/subdir/config3.yml",
}
m.statuses["observability/config1.yml"] = files.FileStatusOK
m.statuses["observability/config2.yml"] = files.FileStatusModified
m.statuses["observability/subdir/config3.yml"] = files.FileStatusMissing
},
expected: map[string]files.FileStatus{
composeFileObservability: files.FileStatusOK,
"observability/config1.yml": files.FileStatusOK,
"observability/config2.yml": files.FileStatusModified,
"observability/subdir/config3.yml": files.FileStatusMissing,
},
},
{
name: "compose_stacks_combined",
stack: ProductStackCompose,
setup: func(m *mockFiles) {
// pentagi
m.statuses[composeFilePentagi] = files.FileStatusOK
// graphiti
m.statuses[composeFileGraphiti] = files.FileStatusOK
// langfuse
m.statuses[composeFileLangfuse] = files.FileStatusModified
// observability
m.statuses[composeFileObservability] = files.FileStatusMissing
m.lists[observabilityDirectory] = []string{
"observability/config.yml",
}
m.statuses["observability/config.yml"] = files.FileStatusOK
},
expected: map[string]files.FileStatus{
composeFilePentagi: files.FileStatusOK,
composeFileGraphiti: files.FileStatusOK,
composeFileLangfuse: files.FileStatusModified,
composeFileObservability: files.FileStatusMissing,
"observability/config.yml": files.FileStatusOK,
},
},
{
name: "all_stacks_combined",
stack: ProductStackAll,
setup: func(m *mockFiles) {
// pentagi
m.statuses[composeFilePentagi] = files.FileStatusOK
// graphiti
m.statuses[composeFileGraphiti] = files.FileStatusOK
// langfuse
m.statuses[composeFileLangfuse] = files.FileStatusModified
// observability
m.statuses[composeFileObservability] = files.FileStatusMissing
m.lists[observabilityDirectory] = []string{
"observability/config.yml",
}
m.statuses["observability/config.yml"] = files.FileStatusOK
},
expected: map[string]files.FileStatus{
composeFilePentagi: files.FileStatusOK,
composeFileGraphiti: files.FileStatusOK,
composeFileLangfuse: files.FileStatusModified,
composeFileObservability: files.FileStatusMissing,
"observability/config.yml": files.FileStatusOK,
},
},
{
name: "unsupported_stack",
stack: ProductStackWorker,
setup: func(m *mockFiles) {},
expected: map[string]files.FileStatus{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mockFiles := newMockFiles()
tt.setup(mockFiles)
// create test processor and fsOps
processor := createTestProcessor()
// set working dir via state env path
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join("/test/working", ".env")
processor.files = mockFiles
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
result, err := fsOps.checkStackIntegrity(t.Context(), tt.stack)
if err != nil {
if tt.stack == ProductStackWorker {
// unsupported stack should return an error or empty result; current impl returns error
return
}
t.Fatalf("unexpected error: %v", err)
}
if len(result) != len(tt.expected) {
t.Errorf("expected %d file statuses, got %d", len(tt.expected), len(result))
}
for path, expectedStatus := range tt.expected {
if actualStatus, ok := result[path]; !ok {
t.Errorf("expected file %s not found in result", path)
} else if actualStatus != expectedStatus {
t.Errorf("file %s: expected status %v, got %v", path, expectedStatus, actualStatus)
}
}
})
}
}
func TestCheckStackIntegrity_RealFiles(t *testing.T) {
// Test with real files interface would require proper embedded files setup
// For now, we focus on mock-based testing which covers the logic
t.Run("mock_based_coverage", func(t *testing.T) {
// The logic is covered by TestGetStackFilesStatus above
// Real files integration would require setting up embedded content
// which is beyond the scope of unit tests
mockFiles := newMockFiles()
// Setup comprehensive test scenario
mockFiles.statuses[composeFilePentagi] = files.FileStatusOK
mockFiles.statuses[composeFileGraphiti] = files.FileStatusOK
mockFiles.statuses[composeFileLangfuse] = files.FileStatusModified
mockFiles.statuses[composeFileObservability] = files.FileStatusMissing
mockFiles.lists[observabilityDirectory] = []string{
"observability/config.yml",
"observability/subdir/nested.yml",
}
mockFiles.statuses["observability/config.yml"] = files.FileStatusOK
mockFiles.statuses["observability/subdir/nested.yml"] = files.FileStatusModified
for _, stack := range []ProductStack{ProductStackAll, ProductStackCompose} {
// create test processor and fsOps
processor := createTestProcessor()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join("/test", ".env")
processor.files = mockFiles
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
result, err := fsOps.checkStackIntegrity(t.Context(), stack)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Verify all files are captured
expectedCount := 6 // 4 compose files + 2 observability directory files
if len(result) != expectedCount {
t.Errorf("expected %d files, got %d", expectedCount, len(result))
}
}
})
}
func TestFileSystemOperations_IntegrityWithForceMode(t *testing.T) {
// Test the interaction between ensure/verify integrity and force mode
t.Run("ensure_respects_force_mode", func(t *testing.T) {
processor := createTestProcessor()
mockFiles := processor.files.(*mockFiles)
// Setup files
mockFiles.statuses[composeFilePentagi] = files.FileStatusModified
mockFiles.AddFile(composeFilePentagi, []byte("embedded content"))
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
// First without force - should not overwrite
state := &operationState{force: false, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.ensureStackIntegrity(t.Context(), ProductStackPentagi, state)
assertNoError(t, err)
// Check that file was not copied (force=false with existing file)
copyCount := 0
for _, copy := range mockFiles.copies {
if copy.Src == composeFilePentagi {
copyCount++
}
}
// Note: in real implementation, existing file check happens in ensureFileFromEmbed
// which uses fileExists check on actual filesystem, not mock
// Now with force - should overwrite
state.force = true
err = fsOps.ensureStackIntegrity(t.Context(), ProductStackPentagi, state)
assertNoError(t, err)
})
}
func TestFileSystemOperations_StackSpecificBehavior(t *testing.T) {
// Test stack-specific behaviors
t.Run("observability_handles_directory", func(t *testing.T) {
processor := createTestProcessor()
mockFiles := processor.files.(*mockFiles)
tmpDir := t.TempDir()
mockState := processor.state.(*mockState)
mockState.envPath = filepath.Join(tmpDir, ".env")
// Setup observability directory
mockFiles.lists[observabilityDirectory] = []string{
"observability/config1.yml",
"observability/nested/config2.yml",
}
mockFiles.statuses["observability/config1.yml"] = files.FileStatusOK
mockFiles.statuses["observability/nested/config2.yml"] = files.FileStatusOK
mockFiles.AddFile(observabilityDirectory, []byte{}) // mark as directory
fsOps := newFileSystemOperations(processor).(*fileSystemOperationsImpl)
state := &operationState{force: false, mx: &sync.Mutex{}, ctx: t.Context()}
err := fsOps.ensureStackIntegrity(t.Context(), ProductStackObservability, state)
assertNoError(t, err)
// Should have attempted to copy both compose file and directory
expectedCopies := 2 // compose file + directory
if len(mockFiles.copies) < expectedCopies {
t.Errorf("expected at least %d copy operations, got %d", expectedCopies, len(mockFiles.copies))
}
})
}
backend/cmd/installer/processor/locale.go
+130
View File
@@ -0,0 +1,130 @@
package processor
// Docker operations messages
const (
MsgPullingImage = "Pulling image: %s"
MsgImagePullCompleted = "Completed pulling %s"
MsgImagePullFailed = "Failed to pull image %s: %v"
MsgRemovingWorkerContainers = "Removing worker containers"
MsgStoppingContainer = "Stopping container %s"
MsgRemovingContainer = "Removing container %s"
MsgContainerRemoved = "Removed container %s"
MsgNoWorkerContainersFound = "No worker containers found"
MsgWorkerContainersRemoved = "Removed %d worker containers"
MsgRemovingImage = "Removing image: %s"
MsgImageRemoved = "Successfully removed image %s"
MsgImageNotFound = "Image %s not found (already removed)"
MsgWorkerImagesRemoveCompleted = "Worker images removal completed"
MsgEnsuringDockerNetworks = "Ensuring docker networks exist"
MsgDockerNetworkExists = "Docker network exists: %s"
MsgCreatingDockerNetwork = "Creating docker network: %s"
MsgDockerNetworkCreated = "Docker network created: %s"
MsgDockerNetworkCreateFailed = "Failed to create docker network %s: %v"
MsgRecreatingDockerNetwork = "Recreating docker network with compose labels: %s"
MsgDockerNetworkRemoved = "Docker network removed: %s"
MsgDockerNetworkRemoveFailed = "Failed to remove docker network %s: %v"
MsgDockerNetworkInUse = "Docker network %s is in use by containers; cannot recreate"
)
// File system operations messages
const (
MsgExtractingDockerCompose = "Extracting docker-compose.yml"
MsgExtractingLangfuseCompose = "Extracting docker-compose-langfuse.yml"
MsgExtractingObservabilityCompose = "Extracting docker-compose-observability.yml"
MsgExtractingObservabilityDirectory = "Extracting observability directory"
MsgSkippingExternalLangfuse = "Skipping external Langfuse deployment"
MsgSkippingExternalObservability = "Skipping external Observability deployment"
MsgPatchingComposeFile = "Patching docker-compose file: %s"
MsgComposePatchCompleted = "Docker-compose file patching completed"
MsgCleaningUpStackFiles = "Cleaning up stack files for %s"
MsgStackFilesCleanupCompleted = "Stack files cleanup completed"
MsgEnsurngStackIntegrity = "Ensuring %s stack integrity"
MsgVerifyingStackIntegrity = "Verifying %s stack integrity"
MsgStackIntegrityVerified = "Stack %s integrity verified"
MsgUpdatingExistingFile = "Updating existing file: %s"
MsgCreatingMissingFile = "Creating missing file: %s"
MsgFileIntegrityValid = "File integrity valid: %s"
MsgSkippingModifiedFile = "Skipping modified files: %s"
MsgDirectoryCheckedWithModified = "Directory checked with modified files present: %s"
)
// Update operations messages
const (
MsgCheckingUpdates = "Checking for updates"
MsgDownloadingInstaller = "Downloading installer update"
MsgInstallerDownloadCompleted = "Installer download completed"
MsgUpdatingInstaller = "Updating installer"
MsgRemovingInstaller = "Removing installer"
MsgInstallerUpdateCompleted = "Installer update completed"
MsgVerifyingBinaryChecksum = "Verifying binary checksum"
MsgReplacingInstallerBinary = "Replacing installer binary"
)
// Remove operations messages
const (
MsgRemovingStack = "Removing stack: %s"
MsgStackRemovalCompleted = "Stack removal completed for %s"
MsgPurgingStack = "Purging stack: %s"
MsgStackPurgeCompleted = "Stack purge completed for %s"
MsgExecutingDockerCompose = "Executing docker-compose command: %s"
MsgDockerComposeCompleted = "Docker-compose command completed"
MsgFactoryResetStarting = "Starting factory reset"
MsgFactoryResetCompleted = "Factory reset completed"
MsgRestoringDefaultEnv = "Restoring default .env from embedded"
MsgDefaultEnvRestored = "Default .env restored"
)
type Subsystem string
const (
SubsystemDocker Subsystem = "docker"
SubsystemCompose Subsystem = "compose"
SubsystemFileSystem Subsystem = "file-system"
SubsystemUpdate Subsystem = "update"
)
type SubsystemOperationMessage struct {
Enter string
Exit string
Error string
}
var SubsystemOperationMessages = map[Subsystem]map[ProcessorOperation]SubsystemOperationMessage{
SubsystemCompose: {
ProcessorOperationStart: SubsystemOperationMessage{
Enter: "Starting %s compose stack",
Exit: "Compose stack %s was started",
Error: "Failed to start %s compose stack",
},
ProcessorOperationStop: SubsystemOperationMessage{
Enter: "Stopping %s compose stack",
Exit: "Compose stack %s was stopped",
Error: "Failed to stop %s compose stack",
},
ProcessorOperationRestart: SubsystemOperationMessage{
Enter: "Restarting %s compose stack",
Exit: "Compose stack %s was restarted",
Error: "Failed to restart %s compose stack",
},
ProcessorOperationUpdate: SubsystemOperationMessage{
Enter: "Updating %s compose stack",
Exit: "Compose stack %s was updated",
Error: "Failed to update %s compose stack",
},
ProcessorOperationDownload: SubsystemOperationMessage{
Enter: "Downloading %s compose stack",
Exit: "Compose stack %s was downloaded",
Error: "Failed to download %s compose stack",
},
ProcessorOperationRemove: SubsystemOperationMessage{
Enter: "Removing %s compose stack",
Exit: "Compose stack %s was removed",
Error: "Failed to remove %s compose stack",
},
ProcessorOperationPurge: SubsystemOperationMessage{
Enter: "Purging %s compose stack",
Exit: "Compose stack %s was purged",
Error: "Failed to purge %s compose stack",
},
},
}
backend/cmd/installer/processor/logic.go
+900
View File
@@ -0,0 +1,900 @@
package processor
import (
"context"
"fmt"
"maps"
"os"
"os/exec"
"path/filepath"
"runtime"
"sync"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/wizard/logger"
)
const (
minTerminalWidthForCompose = 56
)
// helper to validate stack applicability for operation
func (p *processor) validateOperation(stack ProductStack, operation ProcessorOperation) error {
switch operation {
case ProcessorOperationApplyChanges, ProcessorOperationInstall, ProcessorOperationFactoryReset:
if stack != ProductStackAll {
return fmt.Errorf("operation %s not applicable for stack %s", operation, stack)
}
case ProcessorOperationUpdate, ProcessorOperationDownload, ProcessorOperationRemove, ProcessorOperationPurge:
break // can be applied to any stack
case ProcessorOperationStart, ProcessorOperationStop, ProcessorOperationRestart, ProcessorOperationCheckFiles:
if stack == ProductStackWorker || stack == ProductStackInstaller {
return fmt.Errorf("operation %s not applicable for stack %s", operation, stack)
}
case ProcessorOperationResetPassword:
if stack != ProductStackPentagi {
return fmt.Errorf("operation %s only applicable for PentAGI stack", operation)
}
}
return nil
}
// isEmbeddedDeployment checks if the deployment mode is embedded based on environment variable
func (p *processor) isEmbeddedDeployment(stack ProductStack) bool {
switch stack {
case ProductStackObservability:
envVar, envVarValueEmbedded := "OTEL_HOST", checker.DefaultObservabilityEndpoint
if envVar, exists := p.state.GetVar(envVar); exists && envVar.Value == envVarValueEmbedded {
return true
}
return false
case ProductStackLangfuse:
if !p.checker.LangfuseConnected {
return false
}
envVar, envVarValueEmbedded := "LANGFUSE_BASE_URL", checker.DefaultLangfuseEndpoint
if envVar, exists := p.state.GetVar(envVar); exists && envVar.Value == envVarValueEmbedded {
return true
}
return false
case ProductStackGraphiti:
if !p.checker.GraphitiConnected {
return false
}
envVar, envVarValueEmbedded := "GRAPHITI_URL", checker.DefaultGraphitiEndpoint
if envVar, exists := p.state.GetVar(envVar); exists && envVar.Value == envVarValueEmbedded {
return true
}
return false
case ProductStackPentagi, ProductStackWorker, ProductStackInstaller:
return true
default:
return false
}
}
func (p *processor) runCommand(cmd *exec.Cmd, stack ProductStack, state *operationState) error {
if state.terminal != nil {
// patch env vars for docker compose and small size screen
if width, _ := state.terminal.GetSize(); width < minTerminalWidthForCompose || runtime.GOOS == "windows" {
cmd.Env = append(cmd.Env, "COMPOSE_ANSI=never")
}
if err := state.terminal.Execute(cmd); err != nil {
return fmt.Errorf("failed to execute command: %w", err)
}
logger.Log("waiting for command: %s", cmd.String())
if err := cmd.Wait(); err != nil {
return fmt.Errorf("failed to wait for command: %w", err)
}
logger.Log("waiting for terminal to finish: %s", cmd.String())
state.terminal.Wait()
logger.Log("terminal finished: %s", cmd.String())
state.sendOutput(state.terminal.View(), false, stack)
} else {
if output, err := cmd.CombinedOutput(); err != nil {
return fmt.Errorf("failed to run command: %w\n%s", err, string(output))
} else {
state.sendOutput(string(output), true, stack)
}
}
return nil
}
func (p *processor) appendLog(msg string, stack ProductStack, state *operationState) {
if state.terminal != nil {
state.terminal.Append(msg)
state.sendOutput(state.terminal.View(), false, stack)
} else {
state.sendOutput(msg, true, stack)
}
}
func (p *processor) isFileExists(path string) error {
if info, err := os.Stat(path); os.IsNotExist(err) {
return fmt.Errorf("file %s does not exist", path)
} else if err != nil {
return fmt.Errorf("file %s: %w", path, err)
} else if info.IsDir() {
return fmt.Errorf("file %s is a directory", path)
}
return nil
}
func (p *processor) applyChanges(ctx context.Context, state *operationState) (err error) {
stack := ProductStackAll
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
defer func() {
if err != nil {
return
}
// refresh state for updates
if gatherErr := p.checker.GatherUpdatesInfo(ctx); gatherErr != nil {
err = fmt.Errorf("failed to gather info after update: %w", gatherErr)
}
}()
if err := p.validateOperation(stack, ProcessorOperationApplyChanges); err != nil {
return err
}
if !p.state.IsDirty() {
return nil
}
if err = p.state.Commit(); err != nil {
return fmt.Errorf("failed to commit state: %w", err)
}
// refresh checker state after commit to use updated .env values
if err := p.checker.GatherAllInfo(ctx); err != nil {
return fmt.Errorf("failed to gather updated info: %w", err)
}
// ensure required docker networks exist before bringing up stacks
if err := p.dockerOps.ensureMainDockerNetworks(ctx, state); err != nil {
return fmt.Errorf("failed to ensure docker networks: %w", err)
}
// phase 1: Observability Stack Management
if err := p.applyObservabilityChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply observability changes: %w", err)
}
// phase 2: Langfuse Stack Management
if err := p.applyLangfuseChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply langfuse changes: %w", err)
}
// phase 3: Graphiti Stack Management
if err := p.applyGraphitiChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply graphiti changes: %w", err)
}
// phase 4: PentAGI Stack Management (always embedded, always required)
if err := p.applyPentagiChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply pentagi changes: %w", err)
}
return nil
}
func (p *processor) applyObservabilityChanges(ctx context.Context, state *operationState) error {
if p.isEmbeddedDeployment(ProductStackObservability) {
// user wants embedded observability
if !p.checker.ObservabilityExtracted {
// fresh installation - extract all files
if err := p.fsOps.ensureStackIntegrity(ctx, ProductStackObservability, state); err != nil {
return fmt.Errorf("failed to ensure observability integrity: %w", err)
}
} else {
// files exist - verify integrity, update if force=true
if err := p.fsOps.verifyStackIntegrity(ctx, ProductStackObservability, state); err != nil {
return fmt.Errorf("failed to verify observability integrity: %w", err)
}
}
// update/start containers
if err := p.composeOps.updateStack(ctx, ProductStackObservability, state); err != nil {
return fmt.Errorf("failed to update observability stack: %w", err)
}
} else {
// user wants external/disabled observability
if p.checker.ObservabilityInstalled {
// remove containers but keep files (user might re-enable)
if err := p.composeOps.removeStack(ctx, ProductStackObservability, state); err != nil {
return fmt.Errorf("failed to remove observability stack: %w", err)
}
}
}
// refresh state to verify operation success
if err := p.checker.GatherObservabilityInfo(ctx); err != nil {
return fmt.Errorf("failed to gather observability info: %w", err)
}
return nil
}
func (p *processor) applyLangfuseChanges(ctx context.Context, state *operationState) error {
if p.isEmbeddedDeployment(ProductStackLangfuse) {
// user wants embedded langfuse
if !p.checker.LangfuseExtracted {
// fresh installation - extract compose file
if err := p.fsOps.ensureStackIntegrity(ctx, ProductStackLangfuse, state); err != nil {
return fmt.Errorf("failed to ensure langfuse integrity: %w", err)
}
} else {
// file exists - verify integrity, update if force=true
if err := p.fsOps.verifyStackIntegrity(ctx, ProductStackLangfuse, state); err != nil {
return fmt.Errorf("failed to verify langfuse integrity: %w", err)
}
}
// update/start containers
if err := p.composeOps.updateStack(ctx, ProductStackLangfuse, state); err != nil {
return fmt.Errorf("failed to update langfuse stack: %w", err)
}
} else {
// user wants external/disabled langfuse
if p.checker.LangfuseInstalled {
// remove containers but keep files (user might re-enable)
if err := p.composeOps.removeStack(ctx, ProductStackLangfuse, state); err != nil {
return fmt.Errorf("failed to remove langfuse stack: %w", err)
}
}
}
// refresh state to verify operation success
if err := p.checker.GatherLangfuseInfo(ctx); err != nil {
return fmt.Errorf("failed to gather langfuse info: %w", err)
}
return nil
}
func (p *processor) applyGraphitiChanges(ctx context.Context, state *operationState) error {
if p.isEmbeddedDeployment(ProductStackGraphiti) {
// user wants embedded graphiti
if !p.checker.GraphitiExtracted {
// fresh installation - extract compose file
if err := p.fsOps.ensureStackIntegrity(ctx, ProductStackGraphiti, state); err != nil {
return fmt.Errorf("failed to ensure graphiti integrity: %w", err)
}
} else {
// file exists - verify integrity, update if force=true
if err := p.fsOps.verifyStackIntegrity(ctx, ProductStackGraphiti, state); err != nil {
return fmt.Errorf("failed to verify graphiti integrity: %w", err)
}
}
// update/start containers
if err := p.composeOps.updateStack(ctx, ProductStackGraphiti, state); err != nil {
return fmt.Errorf("failed to update graphiti stack: %w", err)
}
} else {
// user wants disabled graphiti
if p.checker.GraphitiInstalled {
// remove containers but keep files (user might re-enable)
if err := p.composeOps.removeStack(ctx, ProductStackGraphiti, state); err != nil {
return fmt.Errorf("failed to remove graphiti stack: %w", err)
}
}
}
// refresh state to verify operation success
if err := p.checker.GatherGraphitiInfo(ctx); err != nil {
return fmt.Errorf("failed to gather graphiti info: %w", err)
}
return nil
}
func (p *processor) applyPentagiChanges(ctx context.Context, state *operationState) error {
// PentAGI is always embedded, always required
if !p.checker.PentagiExtracted {
// fresh installation - extract compose file
if err := p.fsOps.ensureStackIntegrity(ctx, ProductStackPentagi, state); err != nil {
return fmt.Errorf("failed to ensure pentagi integrity: %w", err)
}
} else {
// file exists - verify integrity, update if force=true
if err := p.fsOps.verifyStackIntegrity(ctx, ProductStackPentagi, state); err != nil {
return fmt.Errorf("failed to verify pentagi integrity: %w", err)
}
}
// update/start containers
if err := p.composeOps.updateStack(ctx, ProductStackPentagi, state); err != nil {
return fmt.Errorf("failed to update pentagi stack: %w", err)
}
// refresh state to verify operation success
if err := p.checker.GatherPentagiInfo(ctx); err != nil {
return fmt.Errorf("failed to gather pentagi info: %w", err)
}
return nil
}
// checkFiles computes file statuses for a given stack, honoring the same
// rules as verifyStackIntegrity: active stacks only and excluded files policy.
// It serves as a dry-run for file operations without performing any writes.
func (p *processor) checkFiles(
ctx context.Context, stack ProductStack, state *operationState,
) (result map[string]files.FileStatus, err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
defer func() { state.sendFilesCheck(stack, result, err) }()
result = make(map[string]files.FileStatus)
if err := p.validateOperation(stack, ProcessorOperationCheckFiles); err != nil {
return nil, err
}
switch stack {
case ProductStackPentagi, ProductStackGraphiti, ProductStackLangfuse, ProductStackObservability:
if !p.isEmbeddedDeployment(stack) {
return map[string]files.FileStatus{}, nil
}
result, err = p.fsOps.checkStackIntegrity(ctx, stack)
if err != nil {
return result, fmt.Errorf("failed to check files integrity: %w", err)
}
case ProductStackAll, ProductStackCompose:
for _, s := range allStacks {
if !p.isEmbeddedDeployment(s) {
continue
}
if r, err := p.fsOps.checkStackIntegrity(ctx, s); err != nil {
return result, fmt.Errorf("failed to check %s files integrity: %w", s, err)
} else {
maps.Copy(result, r)
}
}
}
return result, nil
}
func (p *processor) factoryReset(ctx context.Context, state *operationState) (err error) {
stack := ProductStackAll
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
if err := p.validateOperation(stack, ProcessorOperationFactoryReset); err != nil {
return err
}
// step 1: stop and remove stacks with volumes (down -v) with force semantics
p.appendLog(MsgFactoryResetStarting, ProductStackInstaller, state)
if err := p.composeOps.purgeStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to purge stacks: %w", err)
}
// step 2: remove worker containers and volumes in worker env
if err := p.dockerOps.removeWorkerContainers(ctx, state); err != nil {
return fmt.Errorf("failed to remove worker containers: %w", err)
}
if err := p.dockerOps.removeWorkerVolumes(ctx, state); err != nil {
return fmt.Errorf("failed to remove worker volumes: %w", err)
}
// step 3: remove main networks
_ = p.dockerOps.removeMainDockerNetwork(ctx, state, string(ProductDockerNetworkPentagi))
_ = p.dockerOps.removeMainDockerNetwork(ctx, state, string(ProductDockerNetworkObservability))
_ = p.dockerOps.removeMainDockerNetwork(ctx, state, string(ProductDockerNetworkLangfuse))
// step 4: restore .env from embedded and reload state
p.appendLog(MsgRestoringDefaultEnv, ProductStackInstaller, state)
envDir := filepath.Dir(p.state.GetEnvPath())
if err := p.files.Copy(".env", envDir, true); err != nil {
return fmt.Errorf("failed to restore default .env: %w", err)
}
p.appendLog(MsgDefaultEnvRestored, ProductStackInstaller, state)
if err := p.state.Reset(); err != nil {
return fmt.Errorf("failed to reset state: %w", err)
}
// step 5: restore all embedded files to defaults with overwrite
// observability directory and compose files
err = p.fsOps.ensureStackIntegrity(ctx, stack, &operationState{force: true, mx: &sync.Mutex{}, ctx: ctx})
if err != nil {
return fmt.Errorf("failed to restore embedded files: %w", err)
}
p.appendLog(MsgFactoryResetCompleted, ProductStackInstaller, state)
// refresh checker to reflect clean baseline
if err := p.checker.GatherAllInfo(ctx); err != nil {
return fmt.Errorf("failed to gather info after factory reset: %w", err)
}
return nil
}
func (p *processor) install(ctx context.Context, state *operationState) (err error) {
stack := ProductStackAll
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
defer func() {
if err != nil {
return
}
// refresh state for updates
if gatherErr := p.checker.GatherUpdatesInfo(ctx); gatherErr != nil {
err = fmt.Errorf("failed to gather info after update: %w", gatherErr)
}
}()
if err := p.validateOperation(stack, ProcessorOperationInstall); err != nil {
return err
}
// ensure required docker networks exist before bringing up stacks
if err := p.dockerOps.ensureMainDockerNetworks(ctx, state); err != nil {
return fmt.Errorf("failed to ensure docker networks: %w", err)
}
// phase 1: Observability Stack Management
if !p.checker.ObservabilityInstalled {
if err := p.applyObservabilityChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply observability changes: %w", err)
}
}
// phase 2: Langfuse Stack Management
if !p.checker.LangfuseInstalled {
if err := p.applyLangfuseChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply langfuse changes: %w", err)
}
}
// phase 3: Graphiti Stack Management
if !p.checker.GraphitiInstalled {
if err := p.applyGraphitiChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply graphiti changes: %w", err)
}
}
// phase 4: PentAGI Stack Management (always embedded, always required)
if !p.checker.PentagiInstalled {
if err := p.applyPentagiChanges(ctx, state); err != nil {
return fmt.Errorf("failed to apply pentagi changes: %w", err)
}
}
return nil
}
func (p *processor) update(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
defer func() {
if err != nil {
return
}
// refresh state for updates
if gatherErr := p.checker.GatherUpdatesInfo(ctx); gatherErr != nil {
err = fmt.Errorf("failed to gather info after update: %w", gatherErr)
}
}()
allStacks := append(composeOperationAllStacksOrder[ProcessorOperationUpdate],
ProductStackWorker,
ProductStackInstaller,
)
composeStacksUpToDate := map[ProductStack]bool{
ProductStackPentagi: p.checker.PentagiIsUpToDate,
ProductStackGraphiti: p.checker.GraphitiIsUpToDate,
ProductStackLangfuse: p.checker.LangfuseIsUpToDate,
ProductStackObservability: p.checker.ObservabilityIsUpToDate,
}
composeStacksGatherInfo := map[ProductStack]func(ctx context.Context) error{
ProductStackPentagi: p.checker.GatherPentagiInfo,
ProductStackGraphiti: p.checker.GatherGraphitiInfo,
ProductStackLangfuse: p.checker.GatherLangfuseInfo,
ProductStackObservability: p.checker.GatherObservabilityInfo,
}
if err := p.validateOperation(stack, ProcessorOperationUpdate); err != nil {
return err
}
switch stack {
case ProductStackPentagi, ProductStackGraphiti, ProductStackLangfuse, ProductStackObservability:
if composeStacksUpToDate[stack] {
return nil
}
if err := p.composeOps.downloadStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to download stack: %w", err)
}
// docker compose update equivalent for all images
if err := p.composeOps.updateStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to update stack: %w", err)
}
if err := composeStacksGatherInfo[stack](ctx); err != nil {
return fmt.Errorf("failed to gather info after update: %w", err)
}
case ProductStackWorker:
// pull worker images
if err := p.dockerOps.pullWorkerImage(ctx, state); err != nil {
return fmt.Errorf("failed to pull worker images: %w", err)
}
if err := p.checker.GatherWorkerInfo(ctx); err != nil {
return fmt.Errorf("failed to gather worker info after download: %w", err)
}
case ProductStackInstaller:
if p.checker.InstallerIsUpToDate {
return nil
}
if !p.checker.UpdateServerAccessible {
return fmt.Errorf("update server is not accessible")
}
// HTTP GET from update server
return p.updateOps.updateInstaller(ctx, state)
case ProductStackCompose:
for _, s := range composeOperationAllStacksOrder[ProcessorOperationUpdate] {
if err := p.update(ctx, s, state); err != nil {
return err
}
}
case ProductStackAll:
// update all applicable stacks
for _, s := range allStacks {
if err := p.update(ctx, s, state); err != nil {
return err
}
}
default:
return fmt.Errorf("operation update not applicable for stack %s", stack)
}
return nil
}
func (p *processor) download(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
allStacks := append(composeOperationAllStacksOrder[ProcessorOperationDownload],
ProductStackWorker,
ProductStackInstaller,
)
if err := p.validateOperation(stack, ProcessorOperationDownload); err != nil {
return err
}
switch stack {
case ProductStackPentagi, ProductStackGraphiti, ProductStackLangfuse, ProductStackObservability:
// docker compose pull equivalent for all images
if err := p.composeOps.downloadStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to download stack: %w", err)
}
case ProductStackWorker:
// pull worker images
if err := p.dockerOps.pullWorkerImage(ctx, state); err != nil {
return err
}
if err := p.checker.GatherWorkerInfo(ctx); err != nil {
return fmt.Errorf("failed to gather worker info after download: %w", err)
}
if err := p.checker.GatherUpdatesInfo(ctx); err != nil {
return fmt.Errorf("failed to gather worker info after download: %w", err)
}
case ProductStackInstaller:
if p.checker.InstallerIsUpToDate {
return nil
}
if !p.checker.UpdateServerAccessible {
return fmt.Errorf("update server is not accessible")
}
// HTTP GET from update server
return p.updateOps.downloadInstaller(ctx, state)
case ProductStackCompose:
for _, s := range composeOperationAllStacksOrder[ProcessorOperationDownload] {
if err := p.download(ctx, s, state); err != nil {
return err
}
}
case ProductStackAll:
// download all applicable stacks
for _, s := range allStacks {
if err := p.download(ctx, s, state); err != nil {
return err
}
}
default:
return fmt.Errorf("operation download not applicable for stack %s", stack)
}
return nil
}
func (p *processor) remove(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
allStacks := append(composeOperationAllStacksOrder[ProcessorOperationRemove],
ProductStackWorker,
ProductStackInstaller,
)
composeStacksGatherInfo := map[ProductStack]func(ctx context.Context) error{
ProductStackPentagi: p.checker.GatherPentagiInfo,
ProductStackGraphiti: p.checker.GatherGraphitiInfo,
ProductStackLangfuse: p.checker.GatherLangfuseInfo,
ProductStackObservability: p.checker.GatherObservabilityInfo,
}
if err := p.validateOperation(stack, ProcessorOperationRemove); err != nil {
return err
}
switch stack {
case ProductStackPentagi, ProductStackGraphiti, ProductStackLangfuse, ProductStackObservability:
if err := p.composeOps.removeStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to remove stack: %w", err)
}
if err := composeStacksGatherInfo[stack](ctx); err != nil {
return fmt.Errorf("failed to gather info after remove: %w", err)
}
case ProductStackWorker:
// remove worker images and containers
if err := p.dockerOps.removeWorkerImages(ctx, state); err != nil {
return fmt.Errorf("failed to remove worker images: %w", err)
}
if err := p.checker.GatherWorkerInfo(ctx); err != nil {
return fmt.Errorf("failed to gather worker info after remove: %w", err)
}
if err := p.checker.GatherUpdatesInfo(ctx); err != nil {
return fmt.Errorf("failed to gather info after update: %w", err)
}
case ProductStackInstaller:
// remove installer binary
if err := p.updateOps.removeInstaller(ctx, state); err != nil {
return fmt.Errorf("failed to remove installer: %w", err)
}
case ProductStackCompose:
for _, s := range composeOperationAllStacksOrder[ProcessorOperationRemove] {
if err := p.remove(ctx, s, state); err != nil {
return err
}
}
case ProductStackAll:
// remove all stacks
for _, s := range allStacks {
if err := p.remove(ctx, s, state); err != nil {
return err
}
}
default:
return fmt.Errorf("operation remove not applicable for stack %s", stack)
}
return nil
}
func (p *processor) purge(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
defer func() {
if err != nil {
return
}
// refresh state for updates
if gatherErr := p.checker.GatherUpdatesInfo(ctx); gatherErr != nil {
err = fmt.Errorf("failed to gather info after update: %w", gatherErr)
}
}()
allStacks := append(composeOperationAllStacksOrder[ProcessorOperationPurge],
ProductStackWorker,
ProductStackInstaller,
)
composeStacksGatherInfo := map[ProductStack]func(ctx context.Context) error{
ProductStackPentagi: p.checker.GatherPentagiInfo,
ProductStackGraphiti: p.checker.GatherGraphitiInfo,
ProductStackLangfuse: p.checker.GatherLangfuseInfo,
ProductStackObservability: p.checker.GatherObservabilityInfo,
}
if err := p.validateOperation(stack, ProcessorOperationPurge); err != nil {
return err
}
switch stack {
case ProductStackPentagi, ProductStackGraphiti, ProductStackLangfuse, ProductStackObservability:
if err := p.composeOps.purgeImagesStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to purge with images stack: %w", err)
}
if err := composeStacksGatherInfo[stack](ctx); err != nil {
return fmt.Errorf("failed to gather info after purge: %w", err)
}
case ProductStackWorker:
// purge worker images and containers and volumes
if err := p.dockerOps.purgeWorkerImages(ctx, state); err != nil {
return fmt.Errorf("failed to purge worker images: %w", err)
}
if err := p.checker.GatherWorkerInfo(ctx); err != nil {
return fmt.Errorf("failed to gather worker info after purge: %w", err)
}
if err := p.checker.GatherUpdatesInfo(ctx); err != nil {
return fmt.Errorf("failed to gather info after update: %w", err)
}
case ProductStackInstaller:
// remove installer binary
if err := p.updateOps.removeInstaller(ctx, state); err != nil {
return fmt.Errorf("failed to remove installer: %w", err)
}
case ProductStackCompose:
for _, s := range composeOperationAllStacksOrder[ProcessorOperationPurge] {
if err := p.purge(ctx, s, state); err != nil {
return err
}
}
case ProductStackAll:
// purge all stacks
for _, s := range allStacks {
if err := p.purge(ctx, s, state); err != nil {
return err
}
}
// remove custom networks
_ = p.dockerOps.removeMainDockerNetwork(ctx, state, string(ProductDockerNetworkPentagi))
_ = p.dockerOps.removeMainDockerNetwork(ctx, state, string(ProductDockerNetworkObservability))
_ = p.dockerOps.removeMainDockerNetwork(ctx, state, string(ProductDockerNetworkLangfuse))
default:
return fmt.Errorf("operation purge not applicable for stack %s", stack)
}
return nil
}
func (p *processor) start(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
if err := p.validateOperation(stack, ProcessorOperationStart); err != nil {
return err
}
if err := p.composeOps.startStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to start stack: %w", err)
}
if err := p.checker.GatherAllInfo(ctx); err != nil {
return fmt.Errorf("failed to gather info after start: %w", err)
}
return nil
}
func (p *processor) stop(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
if err := p.validateOperation(stack, ProcessorOperationStop); err != nil {
return err
}
if err := p.composeOps.stopStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to stop stack: %w", err)
}
if err := p.checker.GatherAllInfo(ctx); err != nil {
return fmt.Errorf("failed to gather info after stop: %w", err)
}
return nil
}
func (p *processor) restart(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
if err := p.validateOperation(stack, ProcessorOperationRestart); err != nil {
return err
}
if err := p.composeOps.restartStack(ctx, stack, state); err != nil {
return fmt.Errorf("failed to restart stack: %w", err)
}
if err := p.checker.GatherAllInfo(ctx); err != nil {
return fmt.Errorf("failed to gather info after restart: %w", err)
}
return nil
}
func (p *processor) resetPassword(ctx context.Context, stack ProductStack, state *operationState) (err error) {
state.sendStarted(stack)
defer func() { state.sendCompletion(stack, err) }()
if err := p.validateOperation(stack, ProcessorOperationResetPassword); err != nil {
return err
}
if stack != ProductStackPentagi {
return fmt.Errorf("reset password operation only supported for PentAGI stack")
}
if !p.checker.PentagiRunning {
return fmt.Errorf("PentAGI must be running to reset password")
}
if state.passwordValue == "" {
return fmt.Errorf("password value is required")
}
p.appendLog("Resetting admin password...", stack, state)
// perform password reset using PostgreSQL operations
if err := p.performPasswordReset(ctx, state.passwordValue, state); err != nil {
return fmt.Errorf("failed to reset password: %w", err)
}
p.appendLog("Password reset completed successfully", stack, state)
return nil
}
backend/cmd/installer/processor/logic_test.go
+1235
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/processor/mock_test.go
+2176
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/processor/model.go
+230
View File
@@ -0,0 +1,230 @@
package processor
import (
"context"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/state"
"sync"
"time"
tea "github.com/charmbracelet/bubbletea"
)
// processorModel implements interface for bubbletea integration
type processorModel struct {
*processor
}
type ProcessorModel interface {
ApplyChanges(ctx context.Context, opts ...OperationOption) tea.Cmd
CheckFiles(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
FactoryReset(ctx context.Context, opts ...OperationOption) tea.Cmd
Install(ctx context.Context, opts ...OperationOption) tea.Cmd
Update(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
Download(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
Remove(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
Purge(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
Start(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
Stop(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
Restart(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
ResetPassword(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd
HandleMsg(msg tea.Msg) tea.Cmd
}
func NewProcessorModel(state state.State, checker *checker.CheckResult, files files.Files) ProcessorModel {
p := &processor{
mu: &sync.Mutex{},
state: state,
checker: checker,
files: files,
}
// initialize operation handlers with processor instance
p.fsOps = newFileSystemOperations(p)
p.dockerOps = newDockerOperations(p)
p.composeOps = newComposeOperations(p)
p.updateOps = newUpdateOperations(p)
return &processorModel{processor: p}
}
func wrapCommand(
ctx context.Context, stack ProductStack, mu *sync.Mutex, ch <-chan error,
fn func(state *operationState), opts ...OperationOption,
) tea.Cmd {
state := newOperationState(opts)
go func() {
mu.Lock()
fn(state)
mu.Unlock()
}()
teaCmdWaitMsg := func(err error) tea.Cmd {
return func() tea.Msg {
return ProcessorWaitMsg{
ID: state.id,
Error: err,
Operation: state.operation,
Stack: stack,
state: state,
}
}
}
select {
case <-ctx.Done():
return teaCmdWaitMsg(ctx.Err())
case err := <-ch:
return teaCmdWaitMsg(err)
case <-time.After(500 * time.Millisecond):
return teaCmdWaitMsg(nil)
}
}
func (pm *processorModel) ApplyChanges(ctx context.Context, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, ProductStackAll, pm.mu, ch, func(state *operationState) {
ch <- pm.applyChanges(ctx, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationApplyChanges))...)
}
func (pm *processorModel) CheckFiles(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
_, err := pm.checkFiles(ctx, stack, state)
ch <- err
}, append(opts, withContext(ctx), withOperation(ProcessorOperationCheckFiles))...)
}
func (pm *processorModel) FactoryReset(ctx context.Context, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, ProductStackAll, pm.mu, ch, func(state *operationState) {
ch <- pm.factoryReset(ctx, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationFactoryReset))...)
}
func (pm *processorModel) Install(ctx context.Context, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, ProductStackAll, pm.mu, ch, func(state *operationState) {
ch <- pm.install(ctx, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationInstall))...)
}
func (pm *processorModel) Update(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.update(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationUpdate))...)
}
func (pm *processorModel) Download(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.download(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationDownload))...)
}
func (pm *processorModel) Remove(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.remove(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationRemove))...)
}
func (pm *processorModel) Purge(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.purge(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationPurge))...)
}
func (pm *processorModel) Start(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.start(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationStart))...)
}
func (pm *processorModel) Stop(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.stop(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationStop))...)
}
func (pm *processorModel) Restart(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.restart(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationRestart))...)
}
func (pm *processorModel) ResetPassword(ctx context.Context, stack ProductStack, opts ...OperationOption) tea.Cmd {
ch := make(chan error, 1)
return wrapCommand(ctx, stack, pm.mu, ch, func(state *operationState) {
ch <- pm.resetPassword(ctx, stack, state)
}, append(opts, withContext(ctx), withOperation(ProcessorOperationResetPassword))...)
}
func (pm *processorModel) HandleMsg(msg tea.Msg) tea.Cmd {
newWaitMsg := func(num int, stack ProductStack, state *operationState, err error) tea.Cmd {
if state == nil {
return nil // no state, no poll
}
return func() tea.Msg {
return ProcessorWaitMsg{
ID: state.id,
Error: err,
Operation: state.operation,
Stack: stack,
state: state,
num: num,
}
}
}
pollMsg := func(num int, stack ProductStack, state *operationState) tea.Cmd {
if state == nil {
return nil // no state, no poll
}
state.mx.Lock()
ctx := state.ctx
msgs := state.msgs
state.mx.Unlock()
if num < len(msgs) {
nextMsg := msgs[num]
return func() tea.Msg {
return nextMsg
}
}
select {
case <-ctx.Done():
return newWaitMsg(num, stack, state, ctx.Err())
case <-time.After(100 * time.Millisecond):
return newWaitMsg(num, stack, state, nil)
}
}
switch msg := msg.(type) {
case ProcessorWaitMsg:
if msg.Error != nil {
// stop polling after error
return nil
}
return pollMsg(msg.num, msg.Stack, msg.state)
case ProcessorOutputMsg:
return pollMsg(msg.num, msg.Stack, msg.state)
case ProcessorFilesCheckMsg:
return pollMsg(msg.num, msg.Stack, msg.state)
case ProcessorCompletionMsg:
return nil // final message, no poll
case ProcessorStartedMsg:
return pollMsg(msg.num, msg.Stack, msg.state)
default:
return nil // unknown message, no poll
}
}
backend/cmd/installer/processor/pg.go
+93
View File
@@ -0,0 +1,93 @@
package processor
import (
"context"
"database/sql"
"fmt"
_ "github.com/lib/pq"
"golang.org/x/crypto/bcrypt"
)
const (
// PostgreSQL connection constants (fixed for installer on host)
PostgreSQLHost = "127.0.0.1"
PostgreSQLPort = "5432"
// Default values for PostgreSQL configuration
DefaultPostgreSQLUser = "postgres"
DefaultPostgreSQLPassword = "postgres"
DefaultPostgreSQLDatabase = "pentagidb"
// Admin user email
AdminEmail = "admin@pentagi.com"
// Environment variable names
EnvPostgreSQLUser = "PENTAGI_POSTGRES_USER"
EnvPostgreSQLPassword = "PENTAGI_POSTGRES_PASSWORD"
EnvPostgreSQLDatabase = "PENTAGI_POSTGRES_DB"
)
// performPasswordReset updates the admin password in PostgreSQL
func (p *processor) performPasswordReset(ctx context.Context, newPassword string, state *operationState) error {
// get database configuration from state
dbUser := DefaultPostgreSQLUser
if envVar, ok := p.state.GetVar(EnvPostgreSQLUser); ok && envVar.Value != "" {
dbUser = envVar.Value
}
dbPassword := DefaultPostgreSQLPassword
if envVar, ok := p.state.GetVar(EnvPostgreSQLPassword); ok && envVar.Value != "" {
dbPassword = envVar.Value
}
dbName := DefaultPostgreSQLDatabase
if envVar, ok := p.state.GetVar(EnvPostgreSQLDatabase); ok && envVar.Value != "" {
dbName = envVar.Value
}
// create connection string
connStr := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable",
PostgreSQLHost, PostgreSQLPort, dbUser, dbPassword, dbName)
// open database connection
db, err := sql.Open("postgres", connStr)
if err != nil {
return fmt.Errorf("failed to connect to database: %w", err)
}
defer db.Close()
// test connection
if err := db.PingContext(ctx); err != nil {
return fmt.Errorf("failed to ping database: %w", err)
}
p.appendLog(fmt.Sprintf("Connected to PostgreSQL at %s:%s (database: %s)", PostgreSQLHost, PostgreSQLPort, dbName), ProductStackPentagi, state)
// hash the new password
hashedPassword, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
if err != nil {
return fmt.Errorf("failed to hash password: %w", err)
}
// update the admin user password and status
query := `UPDATE users SET password = $1, status = 'active' WHERE mail = $2`
result, err := db.ExecContext(ctx, query, string(hashedPassword), AdminEmail)
if err != nil {
return fmt.Errorf("failed to update password: %w", err)
}
// check if any rows were affected
rowsAffected, err := result.RowsAffected()
if err != nil {
return fmt.Errorf("failed to get rows affected: %w", err)
}
if rowsAffected == 0 {
return fmt.Errorf("no admin user found with email %s", AdminEmail)
}
p.appendLog(fmt.Sprintf("Password updated for %s", AdminEmail), ProductStackPentagi, state)
return nil
}
backend/cmd/installer/processor/processor.go
+253
View File
@@ -0,0 +1,253 @@
package processor
import (
"context"
"sync"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/state"
"pentagi/cmd/installer/wizard/terminal"
)
type ProductStack string
const (
ProductStackPentagi ProductStack = "pentagi"
ProductStackGraphiti ProductStack = "graphiti"
ProductStackLangfuse ProductStack = "langfuse"
ProductStackObservability ProductStack = "observability"
ProductStackCompose ProductStack = "compose"
ProductStackWorker ProductStack = "worker"
ProductStackInstaller ProductStack = "installer"
ProductStackAll ProductStack = "all"
)
type ProcessorOperation string
const (
ProcessorOperationApplyChanges ProcessorOperation = "apply_changes"
ProcessorOperationCheckFiles ProcessorOperation = "check_files"
ProcessorOperationFactoryReset ProcessorOperation = "factory_reset"
ProcessorOperationInstall ProcessorOperation = "install"
ProcessorOperationUpdate ProcessorOperation = "update"
ProcessorOperationDownload ProcessorOperation = "download"
ProcessorOperationRemove ProcessorOperation = "remove"
ProcessorOperationPurge ProcessorOperation = "purge"
ProcessorOperationStart ProcessorOperation = "start"
ProcessorOperationStop ProcessorOperation = "stop"
ProcessorOperationRestart ProcessorOperation = "restart"
ProcessorOperationResetPassword ProcessorOperation = "reset_password"
)
type ProductDockerNetwork string
const (
ProductDockerNetworkPentagi ProductDockerNetwork = "pentagi-network"
ProductDockerNetworkObservability ProductDockerNetwork = "observability-network"
ProductDockerNetworkLangfuse ProductDockerNetwork = "langfuse-network"
)
type FilesCheckResult map[string]files.FileStatus
type Processor interface {
ApplyChanges(ctx context.Context, opts ...OperationOption) error
CheckFiles(ctx context.Context, stack ProductStack, opts ...OperationOption) (FilesCheckResult, error)
FactoryReset(ctx context.Context, opts ...OperationOption) error
Install(ctx context.Context, opts ...OperationOption) error
Update(ctx context.Context, stack ProductStack, opts ...OperationOption) error
Download(ctx context.Context, stack ProductStack, opts ...OperationOption) error
Remove(ctx context.Context, stack ProductStack, opts ...OperationOption) error
Purge(ctx context.Context, stack ProductStack, opts ...OperationOption) error
Start(ctx context.Context, stack ProductStack, opts ...OperationOption) error
Stop(ctx context.Context, stack ProductStack, opts ...OperationOption) error
Restart(ctx context.Context, stack ProductStack, opts ...OperationOption) error
ResetPassword(ctx context.Context, stack ProductStack, opts ...OperationOption) error
}
// WithForce skips validation checks and attempts maximum operations
func WithForce() OperationOption {
return func(c *operationState) { c.force = true }
}
// WithTerminalModel enables embedded terminal model integration
func WithTerminal(term terminal.Terminal) OperationOption {
return func(c *operationState) {
if term != nil {
c.terminal = term
}
}
}
// WithPasswordValue sets password value for reset password operation
func WithPasswordValue(password string) OperationOption {
return func(c *operationState) {
c.passwordValue = password
}
}
// internal interfaces for specialized operations
type fileSystemOperations interface {
ensureStackIntegrity(ctx context.Context, stack ProductStack, state *operationState) error
verifyStackIntegrity(ctx context.Context, stack ProductStack, state *operationState) error
cleanupStackFiles(ctx context.Context, stack ProductStack, state *operationState) error
checkStackIntegrity(ctx context.Context, stack ProductStack) (FilesCheckResult, error)
}
type dockerOperations interface {
pullWorkerImage(ctx context.Context, state *operationState) error
pullDefaultImage(ctx context.Context, state *operationState) error
removeWorkerContainers(ctx context.Context, state *operationState) error
removeWorkerImages(ctx context.Context, state *operationState) error
purgeWorkerImages(ctx context.Context, state *operationState) error
ensureMainDockerNetworks(ctx context.Context, state *operationState) error
removeMainDockerNetwork(ctx context.Context, state *operationState, name string) error
removeMainImages(ctx context.Context, state *operationState, images []string) error
removeWorkerVolumes(ctx context.Context, state *operationState) error
}
type composeOperations interface {
startStack(ctx context.Context, stack ProductStack, state *operationState) error
stopStack(ctx context.Context, stack ProductStack, state *operationState) error
restartStack(ctx context.Context, stack ProductStack, state *operationState) error
updateStack(ctx context.Context, stack ProductStack, state *operationState) error
downloadStack(ctx context.Context, stack ProductStack, state *operationState) error
removeStack(ctx context.Context, stack ProductStack, state *operationState) error
purgeStack(ctx context.Context, stack ProductStack, state *operationState) error
purgeImagesStack(ctx context.Context, stack ProductStack, state *operationState) error
performStackCommand(ctx context.Context, stack ProductStack, state *operationState, args ...string) error
determineComposeFile(stack ProductStack) (string, error)
}
type updateOperations interface {
checkUpdates(ctx context.Context, state *operationState) (*checker.CheckUpdatesResponse, error)
downloadInstaller(ctx context.Context, state *operationState) error
updateInstaller(ctx context.Context, state *operationState) error
removeInstaller(ctx context.Context, state *operationState) error
}
type processor struct {
mu *sync.Mutex
state state.State
checker *checker.CheckResult
files files.Files
// internal operation handlers
fsOps fileSystemOperations
dockerOps dockerOperations
composeOps composeOperations
updateOps updateOperations
}
func NewProcessor(state state.State, checker *checker.CheckResult, files files.Files) Processor {
p := &processor{
mu: &sync.Mutex{},
state: state,
checker: checker,
files: files,
}
// initialize operation handlers with processor instance
p.fsOps = newFileSystemOperations(p)
p.dockerOps = newDockerOperations(p)
p.composeOps = newComposeOperations(p)
p.updateOps = newUpdateOperations(p)
return p
}
func (p *processor) ApplyChanges(ctx context.Context, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationApplyChanges))
return p.applyChanges(ctx, newOperationState(opts))
}
func (p *processor) CheckFiles(ctx context.Context, stack ProductStack, opts ...OperationOption) (FilesCheckResult, error) {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationCheckFiles))
return p.checkFiles(ctx, stack, newOperationState(opts))
}
func (p *processor) FactoryReset(ctx context.Context, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationFactoryReset))
return p.factoryReset(ctx, newOperationState(opts))
}
func (p *processor) Install(ctx context.Context, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationInstall))
return p.install(ctx, newOperationState(opts))
}
func (p *processor) Update(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationUpdate))
return p.update(ctx, stack, newOperationState(opts))
}
func (p *processor) Download(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationDownload))
return p.download(ctx, stack, newOperationState(opts))
}
func (p *processor) Remove(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationRemove))
return p.remove(ctx, stack, newOperationState(opts))
}
func (p *processor) Purge(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationPurge))
return p.purge(ctx, stack, newOperationState(opts))
}
func (p *processor) Start(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationStart))
return p.start(ctx, stack, newOperationState(opts))
}
func (p *processor) Stop(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationStop))
return p.stop(ctx, stack, newOperationState(opts))
}
func (p *processor) Restart(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationRestart))
return p.restart(ctx, stack, newOperationState(opts))
}
func (p *processor) ResetPassword(ctx context.Context, stack ProductStack, opts ...OperationOption) error {
p.mu.Lock()
defer p.mu.Unlock()
opts = append(opts, withContext(ctx), withOperation(ProcessorOperationResetPassword))
return p.resetPassword(ctx, stack, newOperationState(opts))
}
backend/cmd/installer/processor/state.go
+193
View File
@@ -0,0 +1,193 @@
package processor
import (
"context"
"strings"
"sync"
"pentagi/cmd/installer/wizard/terminal"
tea "github.com/charmbracelet/bubbletea"
"github.com/google/uuid"
)
// operationState holds execution options and state for processor operations
type operationState struct {
id string // unique identifier for the command
force bool // attempt maximum operations
terminal terminal.Terminal // embedded terminal model for interactive display
operation ProcessorOperation
passwordValue string // password value for reset password operation
// message chain for reply
mx *sync.Mutex
ctx context.Context
output strings.Builder
msgs []tea.Msg
}
// ProcessorOutputMsg contains command output line
type ProcessorOutputMsg struct {
ID string
Output string
Operation ProcessorOperation
Stack ProductStack
// keeps for continuing the message chain
state *operationState
num int
}
// ProcessorCompletionMsg signals operation completion
type ProcessorCompletionMsg struct {
ID string
Error error
Operation ProcessorOperation
Stack ProductStack
// keeps for continuing the message chain
state *operationState
num int
}
// ProcessorStartedMsg signals operation start
type ProcessorStartedMsg struct {
ID string
Operation ProcessorOperation
Stack ProductStack
// keeps for continuing the message chain
state *operationState
num int
}
// ProcessorWaitMsg signals operation wait for
type ProcessorWaitMsg struct {
ID string
Error error
Operation ProcessorOperation
Stack ProductStack
// keeps for continuing the message chain
state *operationState
num int
}
// ProcessorFilesCheckMsg carries file statuses computed in check
type ProcessorFilesCheckMsg struct {
ID string
Stack ProductStack
Result FilesCheckResult
Error error
// keeps for continuing the message chain
state *operationState
num int
}
type OperationOption func(c *operationState)
func withID(id string) OperationOption {
return func(c *operationState) { c.id = id }
}
func withOperation(operation ProcessorOperation) OperationOption {
return func(c *operationState) { c.operation = operation }
}
func withContext(ctx context.Context) OperationOption {
return func(c *operationState) { c.ctx = ctx }
}
// helper to build operation state with defaults
func newOperationState(opts []OperationOption) *operationState {
state := &operationState{
id: uuid.New().String(),
mx: &sync.Mutex{},
ctx: context.Background(),
msgs: []tea.Msg{},
}
for _, opt := range opts {
opt(state)
}
if state.terminal == nil {
state.terminal = terminal.NewTerminal(
80, 24,
terminal.WithAutoScroll(),
terminal.WithAutoPoll(),
terminal.WithCurrentEnv(),
terminal.WithNoPty(),
)
}
return state
}
// helper to send output message
func (state *operationState) sendOutput(output string, isPartial bool, stack ProductStack) {
state.mx.Lock()
defer state.mx.Unlock()
if isPartial {
state.output.WriteString(output)
state.output.WriteString("\n")
} else {
state.output.Reset()
state.output.WriteString(output)
}
state.msgs = append(state.msgs, ProcessorOutputMsg{
ID: state.id,
Output: state.output.String(),
Operation: state.operation,
Stack: stack,
state: state,
num: len(state.msgs) + 1,
})
}
// helper to send completion message
func (state *operationState) sendCompletion(stack ProductStack, err error) {
state.mx.Lock()
defer state.mx.Unlock()
state.msgs = append(state.msgs, ProcessorCompletionMsg{
ID: state.id,
Error: err,
Operation: state.operation,
Stack: stack,
state: state,
num: len(state.msgs) + 1,
})
}
// helper to send started message
func (state *operationState) sendStarted(stack ProductStack) {
state.mx.Lock()
defer state.mx.Unlock()
state.msgs = append(state.msgs, ProcessorStartedMsg{
ID: state.id,
Operation: state.operation,
Stack: stack,
state: state,
num: len(state.msgs) + 1,
})
}
// helper to send files check message
func (state *operationState) sendFilesCheck(stack ProductStack, result FilesCheckResult, err error) {
state.mx.Lock()
defer state.mx.Unlock()
state.msgs = append(state.msgs, ProcessorFilesCheckMsg{
ID: state.id,
Stack: stack,
Result: result,
Error: err,
state: state,
num: len(state.msgs) + 1,
})
}
backend/cmd/installer/processor/update.go
+278
View File
@@ -0,0 +1,278 @@
package processor
import (
"context"
"crypto/sha256"
"fmt"
"io"
"net/http"
"net/url"
"os"
"runtime"
"time"
"pentagi/cmd/installer/checker"
"pentagi/pkg/version"
)
const updateServerURL = "https://update.pentagi.com"
type updateOperationsImpl struct {
processor *processor
}
func newUpdateOperations(p *processor) updateOperations {
return &updateOperationsImpl{processor: p}
}
func (u *updateOperationsImpl) checkUpdates(ctx context.Context, state *operationState) (*checker.CheckUpdatesResponse, error) {
u.processor.appendLog(MsgCheckingUpdates, ProductStackInstaller, state)
request := u.buildUpdateCheckRequest()
serverURL := u.getUpdateServerURL()
proxyURL := u.getProxyURL()
response, err := u.callUpdateServer(ctx, serverURL, proxyURL, request)
if err != nil {
return nil, fmt.Errorf("failed to check updates: %w", err)
}
return response, nil
}
func (u *updateOperationsImpl) downloadInstaller(ctx context.Context, state *operationState) error {
u.processor.appendLog(MsgDownloadingInstaller, ProductStackInstaller, state)
downloadURL, err := u.getInstallerDownloadURL(ctx)
if err != nil {
return err
}
tempFile, err := u.downloadBinaryToTemp(ctx, downloadURL)
if err != nil {
return err
}
defer os.Remove(tempFile)
u.processor.appendLog(MsgVerifyingBinaryChecksum, ProductStackInstaller, state)
if err := u.verifyBinaryChecksum(tempFile); err != nil {
return err
}
// TODO: copy binary to current update directory
u.processor.appendLog(MsgInstallerUpdateCompleted, ProductStackInstaller, state)
return fmt.Errorf("not implemented")
}
func (u *updateOperationsImpl) updateInstaller(ctx context.Context, state *operationState) error {
u.processor.appendLog(MsgUpdatingInstaller, ProductStackInstaller, state)
downloadURL, err := u.getInstallerDownloadURL(ctx)
if err != nil {
return err
}
tempFile, err := u.downloadBinaryToTemp(ctx, downloadURL)
if err != nil {
return err
}
defer os.Remove(tempFile)
u.processor.appendLog(MsgVerifyingBinaryChecksum, ProductStackInstaller, state)
if err := u.verifyBinaryChecksum(tempFile); err != nil {
return err
}
// TODO: replace installer binary after communication with current installer process
u.processor.appendLog(MsgReplacingInstallerBinary, ProductStackInstaller, state)
if err := u.replaceInstallerBinary(tempFile); err != nil {
return err
}
u.processor.appendLog(MsgInstallerUpdateCompleted, ProductStackInstaller, state)
return fmt.Errorf("not implemented")
}
func (u *updateOperationsImpl) removeInstaller(ctx context.Context, state *operationState) error {
u.processor.appendLog(MsgRemovingInstaller, ProductStackInstaller, state)
// TODO: remove installer binary
return fmt.Errorf("not implemented")
}
func (u *updateOperationsImpl) buildUpdateCheckRequest() checker.CheckUpdatesRequest {
currentVersion := version.GetBinaryVersion()
if versionVar, exists := u.processor.state.GetVar("PENTAGI_VERSION"); exists {
currentVersion = versionVar.Value
}
return checker.CheckUpdatesRequest{
InstallerOsType: runtime.GOOS,
InstallerVersion: currentVersion,
GraphitiConnected: u.processor.checker.GraphitiConnected,
GraphitiExternal: u.processor.checker.GraphitiExternal,
GraphitiInstalled: u.processor.checker.GraphitiInstalled,
LangfuseConnected: u.processor.checker.LangfuseConnected,
LangfuseExternal: u.processor.checker.LangfuseExternal,
LangfuseInstalled: u.processor.checker.LangfuseInstalled,
ObservabilityConnected: u.processor.checker.ObservabilityConnected,
ObservabilityExternal: u.processor.checker.ObservabilityExternal,
ObservabilityInstalled: u.processor.checker.ObservabilityInstalled,
}
}
func (u *updateOperationsImpl) callUpdateServer(
ctx context.Context,
serverURL, proxyURL string,
request checker.CheckUpdatesRequest,
) (*checker.CheckUpdatesResponse, error) {
client := &http.Client{
Timeout: 30 * time.Second,
}
if proxyURL != "" {
proxyURLParsed, err := url.Parse(proxyURL)
if err == nil {
client.Transport = &http.Transport{
Proxy: http.ProxyURL(proxyURLParsed),
}
}
}
return u.callExistingUpdateChecker(ctx, serverURL, client, request)
}
func (u *updateOperationsImpl) getInstallerDownloadURL(ctx context.Context) (string, error) {
response, err := u.checkUpdates(ctx, &operationState{})
if err != nil {
return "", err
}
if response.InstallerIsUpToDate {
return "", fmt.Errorf("no update available")
}
return "https://update.pentagi.com/installer", nil
}
func (u *updateOperationsImpl) downloadBinaryToTemp(ctx context.Context, downloadURL string) (string, error) {
client := &http.Client{
Timeout: 300 * time.Second,
}
resp, err := client.Get(downloadURL)
if err != nil {
return "", fmt.Errorf("failed to download binary: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("download failed with status: %s", resp.Status)
}
tempFile, err := os.CreateTemp("", "pentagi-update-*.bin")
if err != nil {
return "", fmt.Errorf("failed to create temp file: %w", err)
}
defer tempFile.Close()
if _, err := io.Copy(tempFile, resp.Body); err != nil {
os.Remove(tempFile.Name())
return "", fmt.Errorf("failed to write downloaded binary: %w", err)
}
return tempFile.Name(), nil
}
func (u *updateOperationsImpl) verifyBinaryChecksum(filePath string) error {
file, err := os.Open(filePath)
if err != nil {
return fmt.Errorf("failed to open binary for verification: %w", err)
}
defer file.Close()
hasher := sha256.New()
if _, err := io.Copy(hasher, file); err != nil {
return fmt.Errorf("failed to calculate checksum: %w", err)
}
return nil
}
func (u *updateOperationsImpl) replaceInstallerBinary(newBinaryPath string) error {
currentBinary, err := os.Executable()
if err != nil {
return fmt.Errorf("failed to get current binary path: %w", err)
}
backupPath := currentBinary + ".backup"
if err := u.copyFile(currentBinary, backupPath); err != nil {
return fmt.Errorf("failed to create backup: %w", err)
}
if err := u.copyFile(newBinaryPath, currentBinary); err != nil {
u.copyFile(backupPath, currentBinary)
return fmt.Errorf("failed to replace binary: %w", err)
}
if err := os.Chmod(currentBinary, 0755); err != nil {
return fmt.Errorf("failed to set executable permissions: %w", err)
}
os.Remove(backupPath)
return nil
}
func (u *updateOperationsImpl) getUpdateServerURL() string {
if serverVar, exists := u.processor.state.GetVar("UPDATE_SERVER_URL"); exists && serverVar.Value != "" {
return serverVar.Value
}
return "https://update.pentagi.com"
}
func (u *updateOperationsImpl) getProxyURL() string {
if proxyVar, exists := u.processor.state.GetVar("HTTP_PROXY"); exists {
return proxyVar.Value
}
return ""
}
func (u *updateOperationsImpl) copyFile(src, dst string) error {
srcFile, err := os.Open(src)
if err != nil {
return err
}
defer srcFile.Close()
dstFile, err := os.Create(dst)
if err != nil {
return err
}
defer dstFile.Close()
if _, err := io.Copy(dstFile, srcFile); err != nil {
return err
}
return dstFile.Sync()
}
func (u *updateOperationsImpl) callExistingUpdateChecker(
ctx context.Context,
url string,
client *http.Client,
request checker.CheckUpdatesRequest,
) (*checker.CheckUpdatesResponse, error) {
return &checker.CheckUpdatesResponse{
InstallerIsUpToDate: true,
PentagiIsUpToDate: true,
GraphitiIsUpToDate: true,
LangfuseIsUpToDate: true,
ObservabilityIsUpToDate: true,
WorkerIsUpToDate: true,
}, nil
}
backend/cmd/installer/state/example_test.go
+225
View File
@@ -0,0 +1,225 @@
package state
import (
"fmt"
"os"
"path/filepath"
)
// Example demonstrates the complete workflow of state management for .env files
func ExampleState_transactionWorkflow() {
// Setup: Create a test .env file
tmpDir, _ := os.MkdirTemp("", "state_example")
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
initialContent := `# Database Configuration
DATABASE_URL=postgres://localhost:5432/olddb
DATABASE_PASSWORD=old_password
# API Configuration
API_HOST=localhost
API_PORT=8080`
os.WriteFile(envPath, []byte(initialContent), 0644)
fmt.Println("=== PentAGI Configuration Manager ===")
fmt.Println("Starting configuration process...")
// Step 1: Initialize state management
state, err := NewState(envPath)
if err != nil {
panic(err)
}
// Step 2: Multi-step configuration process
fmt.Println("\n--- Step 1: Database Configuration ---")
state.SetStack([]string{"configure_database"})
// User makes changes gradually
state.SetVar("DATABASE_URL", "postgres://prod-server:5432/pentagidb")
state.SetVar("DATABASE_PASSWORD", "secure_prod_password")
state.SetVar("DATABASE_POOL_SIZE", "20")
fmt.Printf("Current step: %s\n", state.GetStack()[0])
fmt.Printf("Modified variables: %d\n", countChangedVars(state))
// Step 3: Continue with API configuration
fmt.Println("\n--- Step 2: API Configuration ---")
state.SetStack([]string{"configure_api"})
state.SetVar("API_HOST", "0.0.0.0")
state.SetVar("API_PORT", "443")
state.SetVar("API_SSL_ENABLED", "true")
fmt.Printf("Current step: %s\n", state.GetStack()[0])
fmt.Printf("Total modified variables: %d\n", countChangedVars(state))
// Show current state
fmt.Println("\n--- Current Configuration ---")
showCurrentConfig(state)
// Step 4: User can choose to commit or reset
fmt.Println("\n--- Decision: Commit Changes ---")
// Commit applies all changes to .env file and cleans up state
err = state.Commit()
if err != nil {
panic(err)
}
fmt.Println("Changes committed successfully!")
fmt.Printf("State file exists: %v\n", state.Exists())
// Output:
// === PentAGI Configuration Manager ===
// Starting configuration process...
//
// --- Step 1: Database Configuration ---
// Current step: configure_database
// Modified variables: 3
//
// --- Step 2: API Configuration ---
// Current step: configure_api
// Total modified variables: 6
//
// --- Current Configuration ---
// DATABASE_URL: postgres://prod-server:5432/pentagidb [CHANGED]
// DATABASE_PASSWORD: secure_prod_password [CHANGED]
// DATABASE_POOL_SIZE: 20 [NEW]
// API_HOST: 0.0.0.0 [CHANGED]
// API_PORT: 443 [CHANGED]
// API_SSL_ENABLED: true [NEW]
//
// --- Decision: Commit Changes ---
// Changes committed successfully!
// State file exists: true
}
// Example demonstrates rollback functionality
func ExampleState_rollbackWorkflow() {
tmpDir, _ := os.MkdirTemp("", "rollback_example")
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
originalContent := "IMPORTANT_SETTING=production_value"
os.WriteFile(envPath, []byte(originalContent), 0644)
fmt.Println("=== Configuration Rollback Example ===")
state, _ := NewState(envPath)
// User starts making risky changes
fmt.Println("Making risky changes...")
state.SetStack([]string{"risky_configuration"})
state.SetVar("IMPORTANT_SETTING", "experimental_value")
state.SetVar("DANGEROUS_SETTING", "could_break_system")
fmt.Printf("Changes pending: %d\n", countChangedVars(state))
// User realizes they made a mistake
fmt.Println("Oops! These changes might break the system...")
fmt.Println("Rolling back all changes...")
// Reset discards all changes and preserves original file
err := state.Reset()
if err != nil {
panic(err)
}
fmt.Println("All changes discarded!")
fmt.Printf("State file exists: %v\n", state.Exists())
// Verify original file is unchanged
content, _ := os.ReadFile(envPath)
fmt.Printf("Original file preserved: %s", string(content))
// Output:
// === Configuration Rollback Example ===
// Making risky changes...
// Changes pending: 2
// Oops! These changes might break the system...
// Rolling back all changes...
// All changes discarded!
// State file exists: true
// Original file preserved: IMPORTANT_SETTING=production_value
}
// Example demonstrates persistence across sessions
func ExampleState_persistenceWorkflow() {
tmpDir, _ := os.MkdirTemp("", "persistence_example")
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
fmt.Println("=== Session Persistence Example ===")
// Session 1: User starts configuration
fmt.Println("Session 1: Starting configuration...")
state1, _ := NewState(envPath)
state1.SetStack([]string{"partial_configuration"})
state1.SetVar("VAR1", "modified_value")
state1.SetVar("VAR2", "new_value")
fmt.Printf("Session 1 - Step: %s, Changes: %d\n",
state1.GetStack()[0], countChangedVars(state1))
// Session 1 ends (simulated by network disconnection, etc.)
fmt.Println("Session 1 ended unexpectedly...")
// Session 2: User reconnects and resumes
fmt.Println("\nSession 2: Resuming configuration...")
state2, _ := NewState(envPath) // Automatically loads saved state
fmt.Printf("Session 2 - Restored Step: %s, Changes: %d\n",
state2.GetStack()[0], countChangedVars(state2))
// Continue from where left off
state2.SetStack([]string{"complete_configuration"})
state2.SetVar("VAR3", "final_value")
fmt.Printf("Session 2 - Final Step: %s, Changes: %d\n",
state2.GetStack()[0], countChangedVars(state2))
// Commit when ready
state2.Commit()
fmt.Println("Configuration completed successfully!")
// Output:
// === Session Persistence Example ===
// Session 1: Starting configuration...
// Session 1 - Step: partial_configuration, Changes: 2
// Session 1 ended unexpectedly...
//
// Session 2: Resuming configuration...
// Session 2 - Restored Step: partial_configuration, Changes: 2
// Session 2 - Final Step: complete_configuration, Changes: 3
// Configuration completed successfully!
}
func countChangedVars(state State) int {
count := 0
for _, envVar := range state.GetAllVars() {
if envVar.IsChanged {
count++
}
}
return count
}
func showCurrentConfig(state State) {
// Show specific variables in fixed order for consistent output
vars := []string{"DATABASE_URL", "DATABASE_PASSWORD", "DATABASE_POOL_SIZE",
"API_HOST", "API_PORT", "API_SSL_ENABLED"}
allVars := state.GetAllVars()
for _, name := range vars {
if envVar, exists := allVars[name]; exists && envVar.IsChanged {
status := " [CHANGED]"
if !envVar.IsPresent() {
status = " [NEW]"
}
fmt.Printf("%s: %s%s\n", name, envVar.Value, status)
}
}
}
backend/cmd/installer/state/state.go
+310
View File
@@ -0,0 +1,310 @@
package state
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"sync"
"time"
"pentagi/cmd/installer/loader"
)
const EULAConsentFile = "eula-consent"
type State interface {
Exists() bool
Reset() error
Commit() error
IsDirty() bool
GetEulaConsent() bool
SetEulaConsent() error
SetStack(stack []string) error
GetStack() []string
GetVar(name string) (loader.EnvVar, bool)
SetVar(name, value string) error
ResetVar(name string) error
GetVars(names []string) (map[string]loader.EnvVar, map[string]bool)
SetVars(vars map[string]string) error
ResetVars(names []string) error
GetAllVars() map[string]loader.EnvVar
GetEnvPath() string
}
type stateData struct {
Stack []string `json:"stack"`
Vars map[string]loader.EnvVar `json:"vars"`
}
type state struct {
mx *sync.Mutex
envPath string
statePath string
stateDir string
stack []string
envFile loader.EnvFile
}
func NewState(envPath string) (State, error) {
envFile, err := loader.LoadEnvFile(envPath)
if err != nil {
return nil, err
}
stateDir := filepath.Join(filepath.Dir(envPath), ".state")
if err := os.MkdirAll(stateDir, 0755); err != nil {
return nil, err
}
envFileName := filepath.Base(envPath)
statePath := filepath.Join(stateDir, fmt.Sprintf("%s.state", envFileName))
s := &state{
mx: &sync.Mutex{},
envPath: envPath,
statePath: statePath,
stateDir: stateDir,
envFile: envFile,
}
if info, err := os.Stat(statePath); err == nil && info.IsDir() {
return nil, fmt.Errorf("'%s' is a directory", statePath)
} else if err == nil {
if err := s.loadState(statePath); err != nil {
return nil, err
}
}
return s, nil
}
func (s *state) Exists() bool {
info, err := os.Stat(s.statePath)
if err != nil {
return false
}
return !info.IsDir()
}
func (s *state) Reset() error {
s.mx.Lock()
defer s.mx.Unlock()
return s.resetState()
}
func (s *state) Commit() error {
s.mx.Lock()
defer s.mx.Unlock()
if err := s.envFile.Save(s.envPath); err != nil {
return err
}
return s.resetState()
}
func (s *state) IsDirty() bool {
s.mx.Lock()
defer s.mx.Unlock()
info, err := os.Stat(s.statePath)
if err != nil {
return false
}
if info.IsDir() {
return false
}
for _, envVar := range s.envFile.GetAll() {
if envVar.IsChanged {
return true
}
}
return false
}
func (s *state) GetEulaConsent() bool {
s.mx.Lock()
defer s.mx.Unlock()
consentFile := filepath.Join(s.stateDir, EULAConsentFile)
if _, err := os.Stat(consentFile); os.IsNotExist(err) {
return false
}
return true
}
func (s *state) SetEulaConsent() error {
s.mx.Lock()
defer s.mx.Unlock()
currentTime := time.Now().Format(time.RFC3339)
consentFile := filepath.Join(s.stateDir, EULAConsentFile)
if err := os.WriteFile(consentFile, []byte(currentTime), 0644); err != nil {
return fmt.Errorf("failed to write eula consent file: %w", err)
}
return nil
}
func (s *state) SetStack(stack []string) error {
s.mx.Lock()
defer s.mx.Unlock()
s.stack = stack
return s.flushState()
}
func (s *state) GetStack() []string {
s.mx.Lock()
defer s.mx.Unlock()
return s.stack
}
func (s *state) GetVar(name string) (loader.EnvVar, bool) {
s.mx.Lock()
defer s.mx.Unlock()
return s.envFile.Get(name)
}
func (s *state) SetVar(name, value string) error {
s.mx.Lock()
defer s.mx.Unlock()
s.envFile.Set(name, value)
return s.flushState()
}
func (s *state) ResetVar(name string) error {
return s.ResetVars([]string{name})
}
func (s *state) GetVars(names []string) (map[string]loader.EnvVar, map[string]bool) {
s.mx.Lock()
defer s.mx.Unlock()
result := make(map[string]loader.EnvVar, len(names))
present := make(map[string]bool, len(names))
for _, name := range names {
envVar, ok := s.envFile.Get(name)
result[name] = envVar
present[name] = ok
}
return result, present
}
func (s *state) SetVars(vars map[string]string) error {
s.mx.Lock()
defer s.mx.Unlock()
for name, value := range vars {
s.envFile.Set(name, value)
}
return s.flushState()
}
func (s *state) ResetVars(names []string) error {
s.mx.Lock()
defer s.mx.Unlock()
envFile, err := loader.LoadEnvFile(s.envPath)
if err != nil {
return err
}
for _, name := range names {
// try to keep valuable variables that are not present in the env file
// but have default value and its default value can be used in the future
if envVar, ok := envFile.Get(name); ok && (envVar.IsPresent() || envVar.Default != "") {
s.envFile.Set(name, envVar.Value)
} else {
s.envFile.Del(name)
}
}
return s.flushState()
}
func (s *state) GetAllVars() map[string]loader.EnvVar {
s.mx.Lock()
defer s.mx.Unlock()
return s.envFile.GetAll()
}
func (s *state) GetEnvPath() string {
return s.envPath
}
func (s *state) loadState(stateFile string) error {
file, err := os.Open(stateFile)
if err != nil {
return fmt.Errorf("failed to open state file: %w", err)
}
var data stateData
if err := json.NewDecoder(file).Decode(&data); err != nil {
// if the state file is corrupted, reset it
data.Stack = []string{}
data.Vars = make(map[string]loader.EnvVar)
}
s.stack = data.Stack
s.envFile.SetAll(data.Vars)
return nil
}
func (s *state) flushState() error {
data := stateData{
Stack: s.stack,
Vars: s.envFile.GetAll(),
}
file, err := os.OpenFile(s.statePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
if err != nil {
return fmt.Errorf("failed to create state file: %w", err)
}
defer file.Close()
if err := json.NewEncoder(file).Encode(data); err != nil {
return fmt.Errorf("failed to encode state file: %w", err)
}
return nil
}
func (s *state) resetState() error {
if err := os.Remove(s.statePath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove state file: %w", err)
}
envFile, err := loader.LoadEnvFile(s.envPath)
if err != nil {
return fmt.Errorf("failed to load state after reset: %w", err)
}
s.envFile = envFile
if err := s.flushState(); err != nil {
return fmt.Errorf("failed to flush state after reset: %w", err)
}
return nil
}
backend/cmd/installer/state/state_test.go
+701
View File
@@ -0,0 +1,701 @@
package state
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestNewState(t *testing.T) {
t.Run("create new state from existing env file", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
// Create test .env file
envContent := `VAR1=value1
VAR2=value2
# Comment
VAR3=value3`
err := os.WriteFile(envPath, []byte(envContent), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create new state: %v", err)
}
if state == nil {
t.Fatal("Expected state to be non-nil")
}
// Check that variables were loaded
allVars := state.GetAllVars()
if len(allVars) < 3 {
t.Errorf("Expected at least 3 variables, got %d", len(allVars))
}
if envVar, exists := state.GetVar("VAR1"); !exists {
t.Error("Expected VAR1 to exist")
} else if envVar.Value != "value1" {
t.Errorf("Expected VAR1 value 'value1', got '%s'", envVar.Value)
}
})
t.Run("create state with non-existent env file", func(t *testing.T) {
_, err := NewState("/non/existent/file.env")
if err == nil {
t.Error("Expected error when creating state with non-existent file")
}
})
t.Run("create state with directory instead of file", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
_, err := NewState(tmpDir)
if err == nil {
t.Error("Expected error when creating state with directory")
}
})
}
func TestStateExists(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
// Create test .env file
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
// Initially no state file exists
if state.Exists() {
t.Error("Expected state to not exist initially")
}
// After setting a variable, state should exist
err = state.SetVar("NEW_VAR", "new_value")
if err != nil {
t.Fatalf("Failed to set variable: %v", err)
}
if !state.Exists() {
t.Error("Expected state to exist after setting variable")
}
}
func TestStateStepManagement(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
// Initially no step
if stack := state.GetStack(); len(stack) > 0 {
t.Errorf("Expected empty stack initially, got '%s'", stack)
}
// Set step
err = state.SetStack([]string{"configure_database"})
if err != nil {
t.Fatalf("Failed to set step: %v", err)
}
if stack := state.GetStack(); len(stack) < 1 {
t.Errorf("Expected step 'configure_database', got '%s'", stack)
}
// Append step
err = state.SetStack(append(state.GetStack(), "configure_api"))
if err != nil {
t.Fatalf("Failed to update step: %v", err)
}
if stack := state.GetStack(); len(stack) < 2 {
t.Errorf("Expected step 'configure_api', got '%s'", stack)
}
}
func TestStateVariableManagement(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("EXISTING_VAR=existing_value"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
t.Run("get existing variable", func(t *testing.T) {
envVar, exists := state.GetVar("EXISTING_VAR")
if !exists {
t.Error("Expected EXISTING_VAR to exist")
}
if envVar.Value != "existing_value" {
t.Errorf("Expected value 'existing_value', got '%s'", envVar.Value)
}
})
t.Run("set new variable", func(t *testing.T) {
err := state.SetVar("NEW_VAR", "new_value")
if err != nil {
t.Fatalf("Failed to set new variable: %v", err)
}
envVar, exists := state.GetVar("NEW_VAR")
if !exists {
t.Error("Expected NEW_VAR to exist")
}
if envVar.Value != "new_value" {
t.Errorf("Expected value 'new_value', got '%s'", envVar.Value)
}
if !envVar.IsChanged {
t.Error("Expected IsChanged to be true for new variable")
}
})
t.Run("update existing variable", func(t *testing.T) {
err := state.SetVar("EXISTING_VAR", "updated_value")
if err != nil {
t.Fatalf("Failed to update existing variable: %v", err)
}
envVar, exists := state.GetVar("EXISTING_VAR")
if !exists {
t.Error("Expected EXISTING_VAR to exist")
}
if envVar.Value != "updated_value" {
t.Errorf("Expected value 'updated_value', got '%s'", envVar.Value)
}
if !envVar.IsChanged {
t.Error("Expected IsChanged to be true for updated variable")
}
})
t.Run("get multiple variables", func(t *testing.T) {
names := []string{"EXISTING_VAR", "NEW_VAR", "NON_EXISTENT"}
vars, present := state.GetVars(names)
if len(vars) != 3 {
t.Errorf("Expected 3 variables in result, got %d", len(vars))
}
if len(present) != 3 {
t.Errorf("Expected 3 presence flags, got %d", len(present))
}
if !present["EXISTING_VAR"] {
t.Error("Expected EXISTING_VAR to be present")
}
if !present["NEW_VAR"] {
t.Error("Expected NEW_VAR to be present")
}
if present["NON_EXISTENT"] {
t.Error("Expected NON_EXISTENT to not be present")
}
if vars["EXISTING_VAR"].Value != "updated_value" {
t.Errorf("Expected EXISTING_VAR value 'updated_value', got '%s'", vars["EXISTING_VAR"].Value)
}
if vars["NEW_VAR"].Value != "new_value" {
t.Errorf("Expected NEW_VAR value 'new_value', got '%s'", vars["NEW_VAR"].Value)
}
})
t.Run("set multiple variables", func(t *testing.T) {
vars := map[string]string{
"BATCH_VAR1": "batch_value1",
"BATCH_VAR2": "batch_value2",
"EXISTING_VAR": "batch_updated",
}
err := state.SetVars(vars)
if err != nil {
t.Fatalf("Failed to set multiple variables: %v", err)
}
for name, expectedValue := range vars {
envVar, exists := state.GetVar(name)
if !exists {
t.Errorf("Expected variable %s to exist after SetVars", name)
continue
}
if envVar.Value != expectedValue {
t.Errorf("Variable %s: expected value %s, got %s", name, expectedValue, envVar.Value)
}
if !envVar.IsChanged {
t.Errorf("Variable %s: expected IsChanged to be true", name)
}
}
})
t.Run("reset single variable", func(t *testing.T) {
// First modify a variable
err := state.SetVar("EXISTING_VAR", "modified_again")
if err != nil {
t.Fatalf("Failed to modify variable: %v", err)
}
// Verify it was changed
envVar, exists := state.GetVar("EXISTING_VAR")
if !exists || envVar.Value != "modified_again" {
t.Fatalf("Variable was not modified as expected")
}
// Reset it
err = state.ResetVar("EXISTING_VAR")
if err != nil {
t.Fatalf("Failed to reset variable: %v", err)
}
// Verify it was reset to original value
envVar, exists = state.GetVar("EXISTING_VAR")
if !exists {
t.Error("Expected EXISTING_VAR to exist after reset")
}
if envVar.Value != "existing_value" {
t.Errorf("Expected EXISTING_VAR to be reset to 'existing_value', got '%s'", envVar.Value)
}
})
t.Run("reset multiple variables", func(t *testing.T) {
// Set some variables first
vars := map[string]string{
"RESET_VAR1": "reset_value1",
"RESET_VAR2": "reset_value2",
"EXISTING_VAR": "modified_value",
}
err := state.SetVars(vars)
if err != nil {
t.Fatalf("Failed to set variables: %v", err)
}
// Reset multiple variables
names := []string{"RESET_VAR1", "RESET_VAR2", "EXISTING_VAR"}
err = state.ResetVars(names)
if err != nil {
t.Fatalf("Failed to reset variables: %v", err)
}
// RESET_VAR1 and RESET_VAR2 should be deleted (not in original file)
if _, exists := state.GetVar("RESET_VAR1"); exists {
t.Error("Expected RESET_VAR1 to be deleted after reset")
}
if _, exists := state.GetVar("RESET_VAR2"); exists {
t.Error("Expected RESET_VAR2 to be deleted after reset")
}
// EXISTING_VAR should be reset to original value
envVar, exists := state.GetVar("EXISTING_VAR")
if !exists {
t.Error("Expected EXISTING_VAR to exist after reset")
}
if envVar.Value != "existing_value" {
t.Errorf("Expected EXISTING_VAR to be reset to 'existing_value', got '%s'", envVar.Value)
}
})
t.Run("reset non-existent variable", func(t *testing.T) {
err := state.ResetVar("NON_EXISTENT_VAR")
if err != nil {
t.Errorf("Expected reset of non-existent variable to succeed, got error: %v", err)
}
})
t.Run("get all variables", func(t *testing.T) {
allVars := state.GetAllVars()
if len(allVars) < 2 {
t.Errorf("Expected at least 2 variables, got %d", len(allVars))
}
if _, exists := allVars["EXISTING_VAR"]; !exists {
t.Error("Expected EXISTING_VAR in GetAllVars result")
}
if _, exists := allVars["NEW_VAR"]; !exists {
t.Error("Expected NEW_VAR in GetAllVars result")
}
})
}
func TestStateCommit(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
originalContent := "ORIGINAL_VAR=original_value"
err := os.WriteFile(envPath, []byte(originalContent), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
// Make changes
err = state.SetStack([]string{"testing_commit"})
if err != nil {
t.Fatalf("Failed to set step: %v", err)
}
err = state.SetVar("ORIGINAL_VAR", "modified_value")
if err != nil {
t.Fatalf("Failed to set variable: %v", err)
}
err = state.SetVar("NEW_VAR", "new_value")
if err != nil {
t.Fatalf("Failed to set new variable: %v", err)
}
// Verify state exists
if !state.Exists() {
t.Error("Expected state to exist before commit")
}
// Commit changes
err = state.Commit()
if err != nil {
t.Fatalf("Failed to commit state: %v", err)
}
// Verify state file was reloaded and exists
if !state.Exists() {
t.Error("Expected state to exist after commit")
}
// Verify .env file was updated
content, err := os.ReadFile(envPath)
if err != nil {
t.Fatalf("Failed to read env file after commit: %v", err)
}
contentStr := string(content)
if !containsLine(contentStr, "ORIGINAL_VAR=modified_value") {
t.Error("Expected ORIGINAL_VAR to be updated in env file")
}
if !containsLine(contentStr, "NEW_VAR=new_value") {
t.Error("Expected NEW_VAR to be added to env file")
}
// Verify backup was created
backupDir := filepath.Join(tmpDir, ".bak")
entries, err := os.ReadDir(backupDir)
if err != nil {
t.Fatalf("Failed to read backup directory: %v", err)
}
if len(entries) == 0 {
t.Error("Expected backup file to be created")
}
}
func TestStateReset(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
originalContent := "ORIGINAL_VAR=original_value"
err := os.WriteFile(envPath, []byte(originalContent), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
// Make changes
err = state.SetStack([]string{"testing_reset"})
if err != nil {
t.Fatalf("Failed to set step: %v", err)
}
err = state.SetVar("ORIGINAL_VAR", "modified_value")
if err != nil {
t.Fatalf("Failed to set variable: %v", err)
}
// Verify state exists
if !state.Exists() {
t.Error("Expected state to exist before reset")
}
// Reset state
err = state.Reset()
if err != nil {
t.Fatalf("Failed to reset state: %v", err)
}
// Verify state file was reloaded and exists
if !state.Exists() {
t.Error("Expected state to exist after reset")
}
// Verify .env file was NOT changed
content, err := os.ReadFile(envPath)
if err != nil {
t.Fatalf("Failed to read env file after reset: %v", err)
}
if string(content) != originalContent {
t.Errorf("Expected env file to remain unchanged after reset, got: %s", string(content))
}
}
func TestStatePersistence(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
// Create first state instance and make changes
state1, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create first state: %v", err)
}
err = state1.SetStack([]string{"persistence_test"})
if err != nil {
t.Fatalf("Failed to set step: %v", err)
}
err = state1.SetVar("PERSISTENT_VAR", "persistent_value")
if err != nil {
t.Fatalf("Failed to set variable: %v", err)
}
// Create second state instance (should load saved state)
state2, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create second state: %v", err)
}
// Verify step was restored
if step := state2.GetStack()[0]; step != "persistence_test" {
t.Errorf("Expected step 'persistence_test', got '%s'", step)
}
// Verify variable was restored
envVar, exists := state2.GetVar("PERSISTENT_VAR")
if !exists {
t.Error("Expected PERSISTENT_VAR to exist in restored state")
}
if envVar.Value != "persistent_value" {
t.Errorf("Expected value 'persistent_value', got '%s'", envVar.Value)
}
}
func TestStateErrors(t *testing.T) {
t.Run("state file is directory", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
// Create directory where state file should be
stateDir := filepath.Join(tmpDir, ".state")
statePath := filepath.Join(stateDir, ".env.state")
err = os.MkdirAll(statePath, 0755) // Create directory instead of file
if err != nil {
t.Fatalf("Failed to create state directory: %v", err)
}
_, err = NewState(envPath)
if err == nil {
t.Error("Expected error when state file is directory")
}
})
t.Run("corrupted state file", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
// Create corrupted state file
stateDir := filepath.Join(tmpDir, ".state")
err = os.MkdirAll(stateDir, 0755)
if err != nil {
t.Fatalf("Failed to create state directory: %v", err)
}
statePath := filepath.Join(stateDir, ".env.state")
err = os.WriteFile(statePath, []byte("invalid json content"), 0644)
if err != nil {
t.Fatalf("Failed to create corrupted state file: %v", err)
}
// try to reload original env file if state file is corrupted
state, err := NewState(envPath)
if err != nil {
t.Errorf("Expected reload of original env file when state file is corrupted: %v", err)
} else {
envVar, exist := state.GetVar("VAR1")
if !exist {
t.Error("Expected VAR1 to exist in restored state")
}
if envVar.Value != "value1" {
t.Errorf("Expected value 'value1', got '%s'", envVar.Value)
}
}
})
t.Run("empty state file", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
// Create empty state file
stateDir := filepath.Join(tmpDir, ".state")
err = os.MkdirAll(stateDir, 0755)
if err != nil {
t.Fatalf("Failed to create state directory: %v", err)
}
statePath := filepath.Join(stateDir, ".env.state")
err = os.WriteFile(statePath, []byte(""), 0644)
if err != nil {
t.Fatalf("Failed to create empty state file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Errorf("Expected reload of original env file when state file is empty: %v", err)
} else {
envVar, exist := state.GetVar("VAR1")
if !exist {
t.Error("Expected VAR1 to exist in restored state")
}
if envVar.Value != "value1" {
t.Errorf("Expected value 'value1', got '%s'", envVar.Value)
}
}
})
t.Run("reset non-existent state should succeed", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
// Reset when no state file exists should succeed (idempotent operation)
err = state.Reset()
if err != nil {
t.Errorf("Expected reset of non-existent state to succeed, got error: %v", err)
}
// Multiple resets should also succeed
err = state.Reset()
if err != nil {
t.Errorf("Expected multiple resets to succeed, got error: %v", err)
}
})
t.Run("reset after commit should succeed", func(t *testing.T) {
tmpDir := t.TempDir()
defer os.RemoveAll(tmpDir)
envPath := filepath.Join(tmpDir, ".env")
err := os.WriteFile(envPath, []byte("VAR1=value1"), 0644)
if err != nil {
t.Fatalf("Failed to create test env file: %v", err)
}
state, err := NewState(envPath)
if err != nil {
t.Fatalf("Failed to create state: %v", err)
}
// Make changes
err = state.SetVar("NEW_VAR", "new_value")
if err != nil {
t.Fatalf("Failed to set variable: %v", err)
}
// Commit (which should reset state internally)
err = state.Commit()
if err != nil {
t.Fatalf("Failed to commit: %v", err)
}
// Additional reset should still succeed
err = state.Reset()
if err != nil {
t.Errorf("Expected reset after commit to succeed, got error: %v", err)
}
})
}
func containsLine(content, line string) bool {
lines := strings.Split(content, "\n")
for _, l := range lines {
if strings.TrimSpace(l) == line {
return true
}
}
return false
}
backend/cmd/installer/wizard/app.go
+259
View File
@@ -0,0 +1,259 @@
package wizard
import (
"context"
"fmt"
"pentagi/cmd/installer/checker"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/navigator"
"pentagi/cmd/installer/processor"
"pentagi/cmd/installer/state"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/models"
"pentagi/cmd/installer/wizard/registry"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
const (
BaseHeaderHeight = 2
BaseFooterHeight = 1
MinHeaderHeight = 1
)
// App represents the main wizard application
type App struct {
files files.Files
styles styles.Styles
window window.Window
registry registry.Registry
navigator navigator.Navigator
processor processor.ProcessorModel
controller controller.Controller
currentModel models.BaseScreenModel
hotkeys map[string]string
}
func NewApp(appState state.State, checkResult checker.CheckResult, files files.Files) *App {
styles := styles.New()
window := window.New()
navigator := navigator.NewNavigator(appState, checkResult)
controller := controller.NewController(appState, files, checkResult)
processor := processor.NewProcessorModel(appState, controller.GetChecker(), files)
registry := registry.NewRegistry(controller, styles, window, files, processor)
if len(navigator.GetStack()) == 0 {
navigator.Push(models.WelcomeScreen)
}
app := &App{
files: files,
styles: styles,
window: window,
registry: registry,
navigator: navigator,
processor: processor,
controller: controller,
}
app.initHotkeysLocale()
app.updateScreenMargins()
app.currentModel = registry.GetScreen(navigator.Current())
return app
}
func (app *App) initHotkeysLocale() {
app.hotkeys = map[string]string{
"up|down": locale.NavUpDown,
"left|right": locale.NavLeftRight,
"pgup|pgdown": locale.NavPgUpPgDown,
"home|end": locale.NavHomeEnd,
"enter": locale.NavEnter,
"y|n": locale.NavYn,
"ctrl+c": locale.NavCtrlC,
"ctrl+s": locale.NavCtrlS,
"ctrl+r": locale.NavCtrlR,
"ctrl+h": locale.NavCtrlH,
"tab": locale.NavTab,
}
}
// updateScreenMargins calculates and sets header/footer margins based on current screen
func (app *App) updateScreenMargins() {
app.window.SetHeaderHeight(lipgloss.Height(app.renderHeader()))
app.window.SetFooterHeight(lipgloss.Height(app.renderFooter()))
}
func (app *App) Init() tea.Cmd {
if cmd := app.currentModel.Init(); cmd != nil {
return tea.Batch(cmd, tea.WindowSize())
}
return tea.WindowSize()
}
func (app *App) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
app.window.SetWindowSize(msg.Width, msg.Height)
// update content size
msg.Width, msg.Height = app.window.GetContentSize()
// update margins for current screen (header/footer might change with new size)
app.updateScreenMargins()
// forward the resize message to all screens
return app, app.registry.HandleMsg(msg)
case tea.KeyMsg:
switch msg.String() {
case "ctrl+q":
logger.Log("[App] QUIT")
return app, tea.Quit
case "esc":
logger.Log("[App] ESC: %s", app.navigator.Current())
if app.navigator.Current() != models.WelcomeScreen && app.navigator.CanGoBack() {
// go back to previous screen
targetScreen := app.navigator.Pop()
app.currentModel = app.registry.GetScreen(targetScreen)
logger.Log("[App] ESC: going back to %s", targetScreen)
// update margins for the new screen
app.updateScreenMargins()
// soft initialize the new screen to synchronize with state
return app, app.currentModel.Init()
}
}
case models.NavigationMsg: // massages from screens
if msg.GoBack && app.navigator.CanGoBack() {
app.currentModel = app.registry.GetScreen(app.navigator.Pop())
}
if msg.Target != "" {
app.navigator.Push(msg.Target)
app.currentModel = app.registry.GetScreen(msg.Target)
}
// update margins for the new screen
app.updateScreenMargins()
// soft initialize the new screen to synchronize with state
return app, app.currentModel.Init()
}
return app, app.forwardMsgToCurrentModel(msg)
}
func (app *App) View() string {
if app.currentModel == nil {
return locale.UILoading
}
// all screens have unified header/footer management
header := app.renderHeader()
footer := app.renderFooter()
if !app.window.IsShowHeader() {
header = ""
}
content := app.currentModel.View()
contentArea := app.styles.Content.
Width(app.window.GetContentWidth()).
Height(app.window.GetContentHeight()).
Render(content)
return lipgloss.JoinVertical(lipgloss.Left, header, contentArea, footer)
}
func (app *App) forwardMsgToCurrentModel(msg tea.Msg) tea.Cmd {
if app.currentModel != nil {
model, cmd := app.currentModel.Update(msg)
if newModel := models.RestoreModel(model); newModel != nil {
app.currentModel = newModel
}
return cmd
}
return nil
}
func (app *App) renderHeader() string {
currentScreen := app.navigator.Current()
baseScreen := currentScreen.GetScreen()
windowWidth := app.window.GetWindowWidth()
switch models.ScreenID(baseScreen) {
case models.WelcomeScreen:
return app.styles.RenderASCIILogo(windowWidth)
default:
// other screens use text title
return app.styles.Header.Width(windowWidth).Render(app.getScreenTitle())
}
}
func (app *App) renderFooter() string {
var actions []string
var progressInfo string
// add special progress info for EULA screen
currentScreen := app.navigator.Current()
if currentScreen.GetScreen() == string(models.EULAScreen) {
if eulaModel, ok := app.currentModel.(*models.EULAModel); ok {
_, atEnd, percent := eulaModel.GetScrollInfo()
progressInfo = fmt.Sprintf(locale.EULAProgress, percent)
if atEnd {
progressInfo += locale.EULAProgressComplete
}
actions = append(actions, progressInfo)
}
}
// add navigation actions
if app.navigator.CanGoBack() && currentScreen.GetScreen() != string(models.WelcomeScreen) {
actions = append(actions, locale.NavBack)
}
actions = append(actions, locale.NavExit)
// get hotkeys from current screen model
if app.currentModel != nil {
hotkeys := app.currentModel.GetFormHotKeys()
for _, hotkey := range hotkeys {
if localeHotKey, ok := app.hotkeys[hotkey]; ok {
actions = append(actions, localeHotKey)
}
}
}
return app.styles.RenderFooter(actions, app.window.GetWindowWidth())
}
func (app *App) getScreenTitle() string {
if app.currentModel != nil {
return app.currentModel.GetFormTitle()
}
return locale.WelcomeFormTitle
}
func Run(ctx context.Context, appState state.State, checkResult checker.CheckResult, files files.Files) error {
app := NewApp(appState, checkResult, files)
p := tea.NewProgram(
app,
tea.WithAltScreen(),
tea.WithMouseCellMotion(),
)
if _, err := p.Run(); err != nil { // ignore the return model, use app.currentModel instead
return fmt.Errorf("failed to run installer wizard: %w", err)
}
return nil
}
backend/cmd/installer/wizard/controller/controller.go
+2571
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/wizard/locale/locale.go
+2395
View File
File diff suppressed because it is too large Load Diff
backend/cmd/installer/wizard/logger/logger.go
+127
View File
@@ -0,0 +1,127 @@
package logger
import (
"io"
"os"
"time"
"github.com/sirupsen/logrus"
)
var log *logrus.Logger
func init() {
log = logrus.New()
logFile := "log.json"
if envLogFile, ok := os.LookupEnv("INSTALLER_LOG_FILE"); ok {
logFile = envLogFile
}
out, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
if err != nil {
log.Fatal("Failed to open log file: ", err)
}
log.Out = out
log.Formatter = &logrus.TextFormatter{
ForceColors: true,
DisableQuote: true,
TimestampFormat: time.TimeOnly,
}
}
func Log(message string, args ...any) {
if log == nil {
logrus.Infof(message, args...)
return
}
if len(args) > 0 {
log.Infof(message, args...)
} else {
log.Info(message)
}
}
func Errorf(message string, args ...any) {
if log == nil {
logrus.Errorf(message, args...)
return
}
if len(args) > 0 {
log.Errorf(message, args...)
} else {
log.Error(message)
}
}
func Debugf(message string, args ...any) {
if log == nil {
logrus.Debugf(message, args...)
return
}
if len(args) > 0 {
log.Debugf(message, args...)
} else {
log.Debug(message)
}
}
func Warnf(message string, args ...any) {
if log == nil {
logrus.Warnf(message, args...)
return
}
if len(args) > 0 {
log.Warnf(message, args...)
} else {
log.Warn(message)
}
}
func Fatalf(message string, args ...any) {
if log == nil {
logrus.Fatalf(message, args...)
return
}
if len(args) > 0 {
log.Fatalf(message, args...)
} else {
log.Fatal(message)
}
}
func Panicf(message string, args ...any) {
if log == nil {
logrus.Panicf(message, args...)
return
}
if len(args) > 0 {
log.Panicf(message, args...)
} else {
log.Panic(message)
}
}
func GetLevel() logrus.Level {
if log == nil {
return logrus.GetLevel()
}
return log.GetLevel()
}
func SetLevel(level logrus.Level) {
if log == nil {
logrus.SetLevel(level)
return
}
log.SetLevel(level)
}
func SetOutput(output io.Writer) {
if log == nil {
logrus.SetOutput(output)
return
}
log.SetOutput(output)
}
backend/cmd/installer/wizard/models/ai_agents_settings_form.go
+401
View File
@@ -0,0 +1,401 @@
package models
import (
"fmt"
"strconv"
"strings"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
tea "github.com/charmbracelet/bubbletea"
)
// AIAgentsSettingsFormModel represents the AI agents settings form
type AIAgentsSettingsFormModel struct {
*BaseScreen
}
// NewAIAgentsSettingsFormModel creates a new AI agents settings form model
func NewAIAgentsSettingsFormModel(c controller.Controller, s styles.Styles, w window.Window) *AIAgentsSettingsFormModel {
m := &AIAgentsSettingsFormModel{}
// create base screen with this model as handler (no list handler needed)
m.BaseScreen = NewBaseScreen(c, s, w, m, nil)
return m
}
// BaseScreenHandler interface implementation
func (m *AIAgentsSettingsFormModel) BuildForm() tea.Cmd {
cfg := m.GetController().GetAIAgentsConfig()
fields := []FormField{
m.createBooleanField(
"ask_user",
locale.ToolsAIAgentsSettingHumanInTheLoop,
locale.ToolsAIAgentsSettingHumanInTheLoopDesc,
cfg.HumanInTheLoop,
),
m.createBooleanField(
"assistant_use_agents",
locale.ToolsAIAgentsSettingUseAgents,
locale.ToolsAIAgentsSettingUseAgentsDesc,
cfg.AssistantUseAgents,
),
m.createBooleanField(
"execution_monitor_enabled",
locale.ToolsAIAgentsSettingExecutionMonitor,
locale.ToolsAIAgentsSettingExecutionMonitorDesc,
cfg.ExecutionMonitorEnabled,
),
m.createIntegerField(
"execution_monitor_same_tool_limit",
locale.ToolsAIAgentsSettingSameToolLimit,
locale.ToolsAIAgentsSettingSameToolLimitDesc,
cfg.ExecutionMonitorSameToolLimit,
1,
50,
),
m.createIntegerField(
"execution_monitor_total_tool_limit",
locale.ToolsAIAgentsSettingTotalToolLimit,
locale.ToolsAIAgentsSettingTotalToolLimitDesc,
cfg.ExecutionMonitorTotalToolLimit,
1,
100,
),
m.createIntegerField(
"max_general_agent_tool_calls",
locale.ToolsAIAgentsSettingMaxGeneralToolCalls,
locale.ToolsAIAgentsSettingMaxGeneralToolCallsDesc,
cfg.MaxGeneralAgentToolCalls,
10,
500,
),
m.createIntegerField(
"max_limited_agent_tool_calls",
locale.ToolsAIAgentsSettingMaxLimitedToolCalls,
locale.ToolsAIAgentsSettingMaxLimitedToolCallsDesc,
cfg.MaxLimitedAgentToolCalls,
5,
200,
),
m.createBooleanField(
"agent_planning_step_enabled",
locale.ToolsAIAgentsSettingTaskPlanning,
locale.ToolsAIAgentsSettingTaskPlanningDesc,
cfg.AgentPlanningStepEnabled,
),
}
m.SetFormFields(fields)
return fields[0].Input.Focus()
}
func (m *AIAgentsSettingsFormModel) createBooleanField(key, title, description string, envVar loader.EnvVar) FormField {
input := NewBooleanInput(m.GetStyles(), m.GetWindow(), envVar)
return FormField{
Key: key,
Title: title,
Description: description,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
Suggestions: input.AvailableSuggestions(),
}
}
func (m *AIAgentsSettingsFormModel) createIntegerField(key, title, description string, envVar loader.EnvVar, min, max int) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), envVar)
// set placeholder with range info
if envVar.Default != "" {
input.Placeholder = fmt.Sprintf("%s (%d-%s)", envVar.Default, min, m.formatNumber(max))
} else {
input.Placeholder = fmt.Sprintf("(%d-%s)", min, m.formatNumber(max))
}
return FormField{
Key: key,
Title: title,
Description: description,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
}
}
func (m *AIAgentsSettingsFormModel) validateBooleanField(value, fieldName string) error {
if value != "" && value != "true" && value != "false" {
return fmt.Errorf("invalid boolean value for %s: %s (must be 'true' or 'false')", fieldName, value)
}
return nil
}
func (m *AIAgentsSettingsFormModel) validateIntegerField(value, fieldName string, min, max int) (int, error) {
if value == "" {
return 0, fmt.Errorf("%s cannot be empty", fieldName)
}
intVal, err := strconv.Atoi(value)
if err != nil {
return 0, fmt.Errorf("invalid integer value for %s: %s", fieldName, value)
}
if intVal < min || intVal > max {
return 0, fmt.Errorf("%s must be between %d and %s", fieldName, min, m.formatNumber(max))
}
return intVal, nil
}
func (m *AIAgentsSettingsFormModel) formatNumber(n int) string {
if n >= 1000 {
return fmt.Sprintf("%d,%03d", n/1000, n%1000)
}
return fmt.Sprintf("%d", n)
}
func (m *AIAgentsSettingsFormModel) GetFormTitle() string {
return locale.ToolsAIAgentsSettingsFormTitle
}
func (m *AIAgentsSettingsFormModel) GetFormDescription() string {
return locale.ToolsAIAgentsSettingsFormDescription
}
func (m *AIAgentsSettingsFormModel) GetFormName() string { return locale.ToolsAIAgentsSettingsFormName }
func (m *AIAgentsSettingsFormModel) GetFormSummary() string { return "" }
func (m *AIAgentsSettingsFormModel) GetFormOverview() string {
var sections []string
sections = append(sections, m.styles.Subtitle.Render(locale.ToolsAIAgentsSettingsFormTitle))
sections = append(sections, "")
sections = append(sections, m.styles.Paragraph.Bold(true).Render(locale.ToolsAIAgentsSettingsFormDescription))
sections = append(sections, "")
sections = append(sections, m.styles.Paragraph.Render(locale.ToolsAIAgentsSettingsFormOverview))
return strings.Join(sections, "\n")
}
func (m *AIAgentsSettingsFormModel) GetCurrentConfiguration() string {
sections := []string{m.GetStyles().Subtitle.Render(m.GetFormName())}
cfg := m.GetController().GetAIAgentsConfig()
// helper function for boolean fields
displayBoolean := func(envVar loader.EnvVar, label string) {
val := envVar.Value
if val == "" {
val = envVar.Default
}
if val == "true" {
sections = append(sections, fmt.Sprintf("• %s: %s", label, m.styles.Success.Render(locale.StatusEnabled)))
} else {
sections = append(sections, fmt.Sprintf("• %s: %s", label, m.styles.Warning.Render(locale.StatusDisabled)))
}
}
// helper function for integer fields
displayInteger := func(envVar loader.EnvVar, label string) {
val := envVar.Value
if val == "" {
val = envVar.Default
}
if val != "" {
sections = append(sections, fmt.Sprintf("• %s: %s", label, m.styles.Info.Render(val)))
} else {
sections = append(sections, fmt.Sprintf("• %s: %s", label, m.styles.Warning.Render("not set")))
}
}
// basic settings
displayBoolean(cfg.HumanInTheLoop, locale.ToolsAIAgentsSettingHumanInTheLoop)
displayBoolean(cfg.AssistantUseAgents, locale.ToolsAIAgentsSettingUseAgents)
// execution monitoring
displayBoolean(cfg.ExecutionMonitorEnabled, locale.ToolsAIAgentsSettingExecutionMonitor)
displayInteger(cfg.ExecutionMonitorSameToolLimit, locale.ToolsAIAgentsSettingSameToolLimit)
displayInteger(cfg.ExecutionMonitorTotalToolLimit, locale.ToolsAIAgentsSettingTotalToolLimit)
// tool call limits
displayInteger(cfg.MaxGeneralAgentToolCalls, locale.ToolsAIAgentsSettingMaxGeneralToolCalls)
displayInteger(cfg.MaxLimitedAgentToolCalls, locale.ToolsAIAgentsSettingMaxLimitedToolCalls)
// task planning
displayBoolean(cfg.AgentPlanningStepEnabled, locale.ToolsAIAgentsSettingTaskPlanning)
return strings.Join(sections, "\n")
}
func (m *AIAgentsSettingsFormModel) IsConfigured() bool {
cfg := m.GetController().GetAIAgentsConfig()
return cfg.HumanInTheLoop.IsPresent() || cfg.HumanInTheLoop.IsChanged ||
cfg.AssistantUseAgents.IsPresent() || cfg.AssistantUseAgents.IsChanged ||
cfg.ExecutionMonitorEnabled.IsPresent() || cfg.ExecutionMonitorEnabled.IsChanged ||
cfg.ExecutionMonitorSameToolLimit.IsPresent() || cfg.ExecutionMonitorSameToolLimit.IsChanged ||
cfg.ExecutionMonitorTotalToolLimit.IsPresent() || cfg.ExecutionMonitorTotalToolLimit.IsChanged ||
cfg.MaxGeneralAgentToolCalls.IsPresent() || cfg.MaxGeneralAgentToolCalls.IsChanged ||
cfg.MaxLimitedAgentToolCalls.IsPresent() || cfg.MaxLimitedAgentToolCalls.IsChanged ||
cfg.AgentPlanningStepEnabled.IsPresent() || cfg.AgentPlanningStepEnabled.IsChanged
}
func (m *AIAgentsSettingsFormModel) GetHelpContent() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.ToolsAIAgentsSettingsFormTitle))
sections = append(sections, "")
sections = append(sections, locale.ToolsAIAgentsSettingsHelp)
return strings.Join(sections, "\n")
}
func (m *AIAgentsSettingsFormModel) HandleSave() error {
fields := m.GetFormFields()
if len(fields) != 8 {
return fmt.Errorf("unexpected number of fields: %d", len(fields))
}
cur := m.GetController().GetAIAgentsConfig()
newCfg := &controller.AIAgentsConfig{
HumanInTheLoop: cur.HumanInTheLoop,
AssistantUseAgents: cur.AssistantUseAgents,
ExecutionMonitorEnabled: cur.ExecutionMonitorEnabled,
ExecutionMonitorSameToolLimit: cur.ExecutionMonitorSameToolLimit,
ExecutionMonitorTotalToolLimit: cur.ExecutionMonitorTotalToolLimit,
MaxGeneralAgentToolCalls: cur.MaxGeneralAgentToolCalls,
MaxLimitedAgentToolCalls: cur.MaxLimitedAgentToolCalls,
AgentPlanningStepEnabled: cur.AgentPlanningStepEnabled,
}
// validate and set each field
for i, field := range fields {
value := strings.TrimSpace(field.Input.Value())
switch field.Key {
case "ask_user":
if err := m.validateBooleanField(value, locale.ToolsAIAgentsSettingHumanInTheLoop); err != nil {
return err
}
newCfg.HumanInTheLoop.Value = value
case "assistant_use_agents":
if err := m.validateBooleanField(value, locale.ToolsAIAgentsSettingUseAgents); err != nil {
return err
}
newCfg.AssistantUseAgents.Value = value
case "execution_monitor_enabled":
if err := m.validateBooleanField(value, locale.ToolsAIAgentsSettingExecutionMonitor); err != nil {
return err
}
newCfg.ExecutionMonitorEnabled.Value = value
case "execution_monitor_same_tool_limit":
if val, err := m.validateIntegerField(value, locale.ToolsAIAgentsSettingSameToolLimit, 1, 50); err != nil {
return err
} else {
newCfg.ExecutionMonitorSameToolLimit.Value = strconv.Itoa(val)
}
case "execution_monitor_total_tool_limit":
if val, err := m.validateIntegerField(value, locale.ToolsAIAgentsSettingTotalToolLimit, 1, 100); err != nil {
return err
} else {
newCfg.ExecutionMonitorTotalToolLimit.Value = strconv.Itoa(val)
}
case "max_general_agent_tool_calls":
if val, err := m.validateIntegerField(value, locale.ToolsAIAgentsSettingMaxGeneralToolCalls, 10, 500); err != nil {
return err
} else {
newCfg.MaxGeneralAgentToolCalls.Value = strconv.Itoa(val)
}
case "max_limited_agent_tool_calls":
if val, err := m.validateIntegerField(value, locale.ToolsAIAgentsSettingMaxLimitedToolCalls, 5, 200); err != nil {
return err
} else {
newCfg.MaxLimitedAgentToolCalls.Value = strconv.Itoa(val)
}
case "agent_planning_step_enabled":
if err := m.validateBooleanField(value, locale.ToolsAIAgentsSettingTaskPlanning); err != nil {
return err
}
newCfg.AgentPlanningStepEnabled.Value = value
default:
return fmt.Errorf("unknown field key at index %d: %s", i, field.Key)
}
}
if err := m.GetController().UpdateAIAgentsConfig(newCfg); err != nil {
return fmt.Errorf("error setting config: %v", err)
}
logger.Log("[AIAgentsSettingsFormModel] SAVE: success")
return nil
}
func (m *AIAgentsSettingsFormModel) HandleReset() {
cfg := m.GetController().ResetAIAgentsConfig()
fields := m.GetFormFields()
if len(fields) >= 1 {
fields[0].Input.SetValue(cfg.HumanInTheLoop.Value)
fields[0].Value = fields[0].Input.Value()
}
if len(fields) >= 2 {
fields[1].Input.SetValue(cfg.AssistantUseAgents.Value)
fields[1].Value = fields[1].Input.Value()
}
if len(fields) >= 3 {
fields[2].Input.SetValue(cfg.ExecutionMonitorEnabled.Value)
fields[2].Value = fields[2].Input.Value()
}
if len(fields) >= 4 {
fields[3].Input.SetValue(cfg.ExecutionMonitorSameToolLimit.Value)
fields[3].Value = fields[3].Input.Value()
}
if len(fields) >= 5 {
fields[4].Input.SetValue(cfg.ExecutionMonitorTotalToolLimit.Value)
fields[4].Value = fields[4].Input.Value()
}
if len(fields) >= 6 {
fields[5].Input.SetValue(cfg.MaxGeneralAgentToolCalls.Value)
fields[5].Value = fields[5].Input.Value()
}
if len(fields) >= 7 {
fields[6].Input.SetValue(cfg.MaxLimitedAgentToolCalls.Value)
fields[6].Value = fields[6].Input.Value()
}
if len(fields) >= 8 {
fields[7].Input.SetValue(cfg.AgentPlanningStepEnabled.Value)
fields[7].Value = fields[7].Input.Value()
}
m.SetFormFields(fields)
}
func (m *AIAgentsSettingsFormModel) OnFieldChanged(fieldIndex int, oldValue, newValue string) {}
func (m *AIAgentsSettingsFormModel) GetFormFields() []FormField { return m.fields }
func (m *AIAgentsSettingsFormModel) SetFormFields(fields []FormField) { m.fields = fields }
// Update method - handle screen-specific input
func (m *AIAgentsSettingsFormModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.KeyMsg:
if cmd := m.HandleFieldInput(msg); cmd != nil {
return m, cmd
}
}
return m, m.BaseScreen.Update(msg)
}
// Compile-time interface validation
var _ BaseScreenModel = (*AIAgentsSettingsFormModel)(nil)
var _ BaseScreenHandler = (*AIAgentsSettingsFormModel)(nil)
backend/cmd/installer/wizard/models/apply_changes.go
+531
View File
@@ -0,0 +1,531 @@
package models
import (
"context"
"fmt"
"sort"
"strings"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/processor"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/terminal"
"pentagi/cmd/installer/wizard/window"
tea "github.com/charmbracelet/bubbletea"
)
// ApplyChangesFormModel represents the Apply Changes form
type ApplyChangesFormModel struct {
*BaseScreen
// processor integration
processor processor.ProcessorModel
running bool
// terminal integration
terminal terminal.Terminal
// files access and integrity state
collecting bool
waitingForChoice bool
outdated map[string]files.FileStatus
}
// NewApplyChangesFormModel creates a new Apply Changes form model
func NewApplyChangesFormModel(
c controller.Controller, s styles.Styles, w window.Window, p processor.ProcessorModel,
) *ApplyChangesFormModel {
m := &ApplyChangesFormModel{
processor: p,
}
// create base screen with this model as handler (no list handler needed)
m.BaseScreen = NewBaseScreen(c, s, w, m, nil)
return m
}
// BaseScreenHandler interface implementation
func (m *ApplyChangesFormModel) BuildForm() tea.Cmd {
// no form fields for this screen - it's a display/action screen
m.SetFormFields([]FormField{})
contentWidth, contentHeight := m.getViewportFormSize()
// setup terminal
if m.terminal == nil {
if !m.isVerticalLayout() {
contentWidth -= 2
}
m.terminal = terminal.NewTerminal(contentWidth-2, contentHeight-1,
terminal.WithAutoScroll(),
terminal.WithAutoPoll(),
terminal.WithCurrentEnv(),
)
} else {
m.terminal.Clear()
}
if m.getChangesCount() == 0 {
m.terminal.Append(locale.ApplyChangesNoChanges)
} else {
m.terminal.Append(locale.ApplyChangesNotStarted)
m.terminal.Append("")
m.terminal.Append(locale.ApplyChangesInstructions)
}
// prevent re-initialization on View() calls
if !m.initialized {
m.initialized = true
} else {
return nil
}
// return terminal's init command to start listening for updates
return m.terminal.Init()
}
func (m *ApplyChangesFormModel) GetFormTitle() string {
return locale.ApplyChangesFormTitle
}
func (m *ApplyChangesFormModel) GetFormDescription() string {
return locale.ApplyChangesFormDescription
}
func (m *ApplyChangesFormModel) GetFormName() string {
return locale.ApplyChangesFormName
}
func (m *ApplyChangesFormModel) GetFormSummary() string {
// terminal viewport takes all available space
return ""
}
func (m *ApplyChangesFormModel) GetFormOverview() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.ApplyChangesFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.ApplyChangesFormDescription))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Render(locale.ApplyChangesFormOverview))
return strings.Join(sections, "\n")
}
func (m *ApplyChangesFormModel) GetCurrentConfiguration() string {
var sections []string
config := m.GetController().GetApplyChangesConfig()
sections = append(sections, m.GetStyles().Subtitle.Render(locale.ApplyChangesChangesTitle))
sections = append(sections, "")
if config.ChangesCount == 0 {
sections = append(sections, m.GetStyles().Warning.Render(locale.ApplyChangesChangesEmpty))
return strings.Join(sections, "\n")
}
// show changes count
sections = append(sections, m.GetStyles().Info.Render(
fmt.Sprintf(locale.ApplyChangesChangesCount, config.ChangesCount)))
sections = append(sections, "")
// show warnings if applicable
if config.HasCritical {
sections = append(sections, m.GetStyles().Warning.Render(locale.ApplyChangesWarningCritical))
}
if config.HasSecrets {
sections = append(sections, m.GetStyles().Info.Render(locale.ApplyChangesWarningSecrets))
}
if config.HasCritical || config.HasSecrets {
sections = append(sections, "")
}
// show notes
sections = append(sections, m.GetStyles().Muted.Render(locale.ApplyChangesNoteBackup))
sections = append(sections, m.GetStyles().Muted.Render(locale.ApplyChangesNoteTime))
sections = append(sections, "")
getMaskedValue := func(value string) string {
maskedValue := strings.Repeat("*", len(value))
if len(value) > 15 {
maskedValue = maskedValue[:15] + "..."
} else if len(value) == 0 {
maskedValue = locale.ApplyChangesChangesMasked
}
return maskedValue
}
// list all changes
for _, change := range config.Changes {
value := change.NewValue
if change.Masked {
value = getMaskedValue(value)
}
line := fmt.Sprintf("• %s: %s",
change.Description,
m.GetStyles().Info.Render(value))
sections = append(sections, line)
}
return strings.Join(sections, "\n")
}
func (m *ApplyChangesFormModel) IsConfigured() bool {
return m.getChangesCount() > 0
}
func (m *ApplyChangesFormModel) GetHelpContent() string {
var sections []string
config := m.GetController().GetApplyChangesConfig()
sections = append(sections, m.GetStyles().Subtitle.Render(locale.ApplyChangesHelpTitle))
sections = append(sections, "")
// show installation or update description based on current state
if !config.IsInstalled {
sections = append(sections, locale.ApplyChangesInstallNotFound)
sections = append(sections, "")
// add additional components if selected
if config.LangfuseEnabled {
sections = append(sections, locale.ApplyChangesInstallFoundLangfuse)
}
if config.ObservabilityEnabled {
sections = append(sections, locale.ApplyChangesInstallFoundObservability)
}
} else {
sections = append(sections, locale.ApplyChangesUpdateFound)
}
sections = append(sections, "")
sections = append(sections, locale.ApplyChangesHelpContent)
sections = append(sections, "")
sections = append(sections, m.GetCurrentConfiguration())
return strings.Join(sections, "\n")
}
func (m *ApplyChangesFormModel) HandleSave() error {
// saving is handled by the processor integration
return nil
}
func (m *ApplyChangesFormModel) HandleReset() {
// reset current changes
m.GetController().Reset()
}
func (m *ApplyChangesFormModel) OnFieldChanged(fieldIndex int, oldValue, newValue string) {
// no fields to change in this screen
}
func (m *ApplyChangesFormModel) GetFormFields() []FormField {
return m.BaseScreen.fields
}
func (m *ApplyChangesFormModel) SetFormFields(fields []FormField) {
m.BaseScreen.fields = fields
}
func (m *ApplyChangesFormModel) getChangesCount() int {
if m.GetController().IsDirty() {
return m.GetController().GetApplyChangesConfig().ChangesCount
}
return 0
}
func (m *ApplyChangesFormModel) handleCompletion(msg processor.ProcessorCompletionMsg) {
if msg.Operation != processor.ProcessorOperationApplyChanges {
return
}
m.running = false
if msg.Error != nil {
m.terminal.Append(fmt.Sprintf("%s: %v\n", locale.ApplyChangesFailed, msg.Error))
} else {
switch msg.Operation {
case processor.ProcessorOperationFactoryReset:
m.terminal.Append(locale.ApplyChangesResetCompleted)
case processor.ProcessorOperationApplyChanges:
m.terminal.Append(locale.ApplyChangesCompleted)
}
}
// rebuild display
m.updateViewports()
}
func (m *ApplyChangesFormModel) handleApplyChanges() tea.Cmd {
if m.terminal != nil {
m.terminal.Clear()
m.terminal.Append(locale.ApplyChangesInProgress)
}
return m.processor.ApplyChanges(context.Background(), processor.WithTerminal(m.terminal))
}
func (m *ApplyChangesFormModel) handleResetChanges() tea.Cmd {
if m.terminal != nil {
m.terminal.Clear()
}
if err := m.GetController().Reset(); err != nil {
if m.terminal != nil {
m.terminal.Append(fmt.Sprintf("%s: %v\n", locale.ApplyChangesFailed, err))
}
} else {
if m.terminal != nil {
m.terminal.Append(locale.ApplyChangesResetCompleted)
}
}
m.updateViewports()
return nil
}
// renderLeftPanel renders the terminal output
func (m *ApplyChangesFormModel) renderLeftPanel() string {
if m.terminal != nil {
return m.terminal.View()
}
// fallback if terminal not initialized
return m.GetStyles().Error.Render(locale.ApplyChangesTerminalIsNotInitialized)
}
// Update method - handle screen-specific input and messages
func (m *ApplyChangesFormModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
handleTerminal := func(msg tea.Msg) (tea.Model, tea.Cmd) {
if m.terminal == nil {
return m, nil
}
updatedModel, cmd := m.terminal.Update(msg)
if terminalModel := terminal.RestoreModel(updatedModel); terminalModel != nil {
m.terminal = terminalModel
}
return m, cmd
}
switch msg := msg.(type) {
case tea.WindowSizeMsg:
contentWidth, contentHeight := m.getViewportFormSize()
// update terminal size when window size changes
if m.terminal != nil {
if !m.isVerticalLayout() {
contentWidth -= 2
}
m.terminal.SetSize(contentWidth-2, contentHeight-1)
}
m.updateViewports()
return m, nil
case terminal.TerminalUpdateMsg:
return handleTerminal(msg)
case processor.ProcessorCompletionMsg:
m.handleCompletion(msg)
return m, m.processor.HandleMsg(msg)
case processor.ProcessorStartedMsg:
return m, m.processor.HandleMsg(msg)
case processor.ProcessorOutputMsg:
// ignore (handled by terminal)
return m, m.processor.HandleMsg(msg)
case processor.ProcessorWaitMsg:
return m, m.processor.HandleMsg(msg)
case processor.ProcessorFilesCheckMsg:
// finish collecting and process result
m.collecting = false
if msg.Error != nil {
if m.terminal != nil {
m.terminal.Append("")
m.terminal.Append(fmt.Sprintf("%s: %v", locale.ApplyChangesFailed, msg.Error))
}
m.updateViewports()
return m, m.processor.HandleMsg(msg)
}
// filter only modified files
outdated := map[string]files.FileStatus{}
for path, st := range msg.Result {
if st == files.FileStatusModified {
outdated[path] = st
}
}
m.outdated = outdated
if len(m.outdated) == 0 {
if m.terminal != nil {
m.terminal.Append("")
m.terminal.Append(locale.ApplyChangesIntegrityNoOutdated)
}
m.running = true
m.updateViewports()
return m, m.handleApplyChanges()
}
m.waitingForChoice = true
if m.terminal != nil {
m.terminal.Append("")
m.terminal.Append(locale.ApplyChangesIntegrityPromptMessage)
m.terminal.Append("")
m.terminal.Append(m.renderOutdatedFiles(m.outdated))
m.terminal.Append("")
m.terminal.Append("(y/n)")
}
m.updateViewports()
return m, m.processor.HandleMsg(msg)
case tea.KeyMsg:
if m.terminal != nil && m.terminal.IsRunning() {
return handleTerminal(msg)
}
switch msg.String() {
case "enter":
if !m.running && !m.collecting && !m.waitingForChoice && m.getChangesCount() != 0 {
m.collecting = true
if m.terminal != nil {
m.terminal.Clear()
m.terminal.Append(locale.ApplyChangesIntegrityPromptTitle)
m.terminal.Append("")
m.terminal.Append(locale.ApplyChangesIntegrityChecking)
}
return m, m.collectOutdatedFiles()
}
return m, nil
case "y":
if !m.running && m.waitingForChoice && m.getChangesCount() != 0 {
m.waitingForChoice = false
m.running = true
return m, m.processor.ApplyChanges(context.Background(), processor.WithTerminal(m.terminal), processor.WithForce())
}
return m, nil
case "n":
if !m.running && m.waitingForChoice && m.getChangesCount() != 0 {
m.waitingForChoice = false
m.running = true
return m, m.handleApplyChanges()
}
return m, nil
case "ctrl+c":
if (m.collecting || m.waitingForChoice) && !m.running {
m.collecting = false
m.waitingForChoice = false
m.outdated = nil
if m.terminal != nil {
m.terminal.Clear()
if m.getChangesCount() == 0 {
m.terminal.Append(locale.ApplyChangesNoChanges)
} else {
m.terminal.Append(locale.ApplyChangesNotStarted)
m.terminal.Append("")
m.terminal.Append(locale.ApplyChangesInstructions)
}
}
m.updateViewports()
return m, nil
}
return m, nil
case "ctrl+r":
if !m.running && !m.collecting && !m.waitingForChoice && m.getChangesCount() != 0 {
return m, m.handleResetChanges()
}
return m, nil
}
// then pass other keys to terminal for scrolling etc.
return handleTerminal(msg)
default:
return handleTerminal(msg)
}
}
// Override View to use custom layout
func (m *ApplyChangesFormModel) View() string {
contentWidth, contentHeight := m.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 {
return locale.UILoading
}
if !m.initialized {
m.handler.BuildForm()
m.fields = m.GetFormFields()
m.updateViewports()
}
leftPanel := m.renderLeftPanel()
rightPanel := m.renderHelp()
if m.isVerticalLayout() {
return m.renderVerticalLayout(leftPanel, rightPanel, contentWidth, contentHeight)
}
return m.renderHorizontalLayout(leftPanel, rightPanel, contentWidth, contentHeight)
}
// GetFormHotKeys returns the hotkeys for this screen
func (m *ApplyChangesFormModel) GetFormHotKeys() []string {
var hotkeys []string
if m.terminal != nil && !m.terminal.IsRunning() && m.getChangesCount() != 0 {
if m.collecting {
hotkeys = append(hotkeys, "ctrl+c")
} else if m.waitingForChoice {
hotkeys = append(hotkeys, "y|n")
hotkeys = append(hotkeys, "ctrl+c")
} else if !m.running {
hotkeys = append(hotkeys, "enter")
hotkeys = append(hotkeys, "ctrl+r")
}
}
return hotkeys
}
// integrity helpers
func (m *ApplyChangesFormModel) collectOutdatedFiles() tea.Cmd {
// delegate file check to processor model; messages will be delivered as ProcessorFilesCheckMsg
return m.processor.CheckFiles(context.Background(), processor.ProductStackAll)
}
func (m *ApplyChangesFormModel) renderOutdatedFiles(outdated map[string]files.FileStatus) string {
if len(outdated) == 0 {
return ""
}
var b strings.Builder
list := make([]string, 0, len(outdated))
for path, st := range outdated {
if st == files.FileStatusModified {
list = append(list, path)
}
}
sort.Strings(list)
for _, path := range list {
b.WriteString("• ")
b.WriteString(path)
b.WriteString("\n")
}
return b.String()
}
// Compile-time interface validation
var _ BaseScreenModel = (*ApplyChangesFormModel)(nil)
var _ BaseScreenHandler = (*ApplyChangesFormModel)(nil)
backend/cmd/installer/wizard/models/base_controls.go
+43
View File
@@ -0,0 +1,43 @@
package models
import (
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/textinput"
)
func NewBooleanInput(styles styles.Styles, window window.Window, envVar loader.EnvVar) textinput.Model {
input := textinput.New()
input.Prompt = ""
input.PlaceholderStyle = styles.FormPlaceholder
input.ShowSuggestions = true
input.SetSuggestions([]string{"true", "false"})
if envVar.Default != "" {
input.Placeholder = envVar.Default
}
if envVar.IsPresent() || envVar.IsChanged {
input.SetValue(envVar.Value)
}
return input
}
func NewTextInput(styles styles.Styles, window window.Window, envVar loader.EnvVar) textinput.Model {
input := textinput.New()
input.Prompt = ""
input.PlaceholderStyle = styles.FormPlaceholder
if envVar.Default != "" {
input.Placeholder = envVar.Default
}
if envVar.IsPresent() || envVar.IsChanged {
input.SetValue(envVar.Value)
}
return input
}
backend/cmd/installer/wizard/models/base_screen.go
+978
View File
@@ -0,0 +1,978 @@
package models
import (
"fmt"
"io"
"slices"
"strings"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/list"
"github.com/charmbracelet/bubbles/textinput"
"github.com/charmbracelet/bubbles/viewport"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
var (
verticalLayoutPaddings = []int{0, 4, 0, 2}
horizontalLayoutPaddings = []int{0, 2, 0, 2}
)
// BaseListOption represents a generic list option that can be used in any list
type BaseListOption struct {
Value string // the actual value
Display string // the display text (can be different from value)
}
func (d BaseListOption) FilterValue() string { return d.Value }
// BaseListDelegate handles rendering of generic list options
type BaseListDelegate struct {
style lipgloss.Style
width int
selectedFg lipgloss.Color
normalFg lipgloss.Color
}
// NewBaseListDelegate creates a new generic list delegate
func NewBaseListDelegate(style lipgloss.Style, width int) *BaseListDelegate {
return &BaseListDelegate{
style: style,
width: width,
selectedFg: styles.Primary,
normalFg: lipgloss.Color(""),
}
}
// SetColors allows customizing the colors
func (d *BaseListDelegate) SetColors(selectedFg, normalFg lipgloss.Color) {
d.selectedFg = selectedFg
d.normalFg = normalFg
}
func (d *BaseListDelegate) SetWidth(width int) { d.width = width }
func (d BaseListDelegate) Height() int { return 1 }
func (d BaseListDelegate) Spacing() int { return 0 }
func (d BaseListDelegate) Update(_ tea.Msg, _ *list.Model) tea.Cmd { return nil }
func (d BaseListDelegate) Render(w io.Writer, m list.Model, index int, listItem list.Item) {
option, ok := listItem.(BaseListOption)
if !ok {
return
}
str := option.Display
if index == m.Index() {
str = d.style.Width(d.width).Foreground(d.selectedFg).Render(str)
} else {
str = d.style.Width(d.width).Foreground(d.normalFg).Render(str)
}
fmt.Fprint(w, str)
}
// BaseListHelper provides utility functions for working with lists
type BaseListHelper struct{}
// CreateList creates a new list with the given options and delegate
func (h BaseListHelper) CreateList(options []BaseListOption, delegate list.ItemDelegate, width, height int) list.Model {
items := make([]list.Item, len(options))
for i, option := range options {
items[i] = option
}
listModel := list.New(items, delegate, width, height)
listModel.SetShowStatusBar(false)
listModel.SetFilteringEnabled(false)
listModel.SetShowHelp(false)
listModel.SetShowTitle(false)
return listModel
}
// SelectByValue selects the list item that matches the given value
func (h BaseListHelper) SelectByValue(listModel *list.Model, value string) {
items := listModel.Items()
for i, item := range items {
if option, ok := item.(BaseListOption); ok && option.Value == value {
listModel.Select(i)
break
}
}
}
// GetSelectedValue returns the value of the currently selected item
func (h BaseListHelper) GetSelectedValue(listModel *list.Model) string {
selectedItem := listModel.SelectedItem()
if selectedItem == nil {
return ""
}
if option, ok := selectedItem.(BaseListOption); ok {
return option.Value
}
return ""
}
// GetSelectedDisplay returns the display text of the currently selected item
func (h BaseListHelper) GetSelectedDisplay(listModel *list.Model) string {
selectedItem := listModel.SelectedItem()
if selectedItem == nil {
return ""
}
if option, ok := selectedItem.(BaseListOption); ok {
return option.Display
}
return ""
}
// BaseScreenModel defines methods that concrete screens must implement
type BaseScreenModel interface {
// GetFormTitle returns the title for the form (layout header)
GetFormTitle() string
// GetFormDescription returns the description for the form (right panel)
GetFormDescription() string
// GetFormName returns the name for the form (right panel)
GetFormName() string
// GetFormOverview returns form overview for list screens (right panel)
GetFormOverview() string
// GetCurrentConfiguration returns text with current configuration for the list screens
GetCurrentConfiguration() string
// IsConfigured returns true if the form is configured
IsConfigured() bool
// GetFormHotKeys returns the hotkeys for the form (layout footer)
GetFormHotKeys() []string
tea.Model // for common interface logic
}
// BaseScreenHandler defines methods that concrete screens must implement
type BaseScreenHandler interface {
// BuildForm builds the specific form fields for this screen
BuildForm() tea.Cmd
// GetFormSummary returns optional summary for the form bottom
GetFormSummary() string
// GetHelpContent returns the right panel help content
GetHelpContent() string
// HandleSave handles saving the form data
HandleSave() error
// HandleReset handles resetting the form to default values
HandleReset()
// OnFieldChanged is called when a form field value changes
OnFieldChanged(fieldIndex int, oldValue, newValue string)
// GetFormFields returns the current form fields
GetFormFields() []FormField
// SetFormFields sets the form fields
SetFormFields(fields []FormField)
}
// BaseListHandler defines methods for screens that use lists (optional)
type BaseListHandler interface {
// GetList returns the list model if this screen uses a list
GetList() *list.Model
// GetListDelegate returns the list delegate if this screen uses a list
GetListDelegate() *BaseListDelegate
// OnListSelectionChanged is called when list selection changes
OnListSelectionChanged(oldSelection, newSelection string)
// GetListTitle returns the title of the list
GetListTitle() string
// GetListDescription returns the description of the list
GetListDescription() string
}
// FormField represents a single form field
type FormField struct {
Key string
Title string
Description string
Placeholder string
Required bool
Masked bool
Input textinput.Model
Value string
Suggestions []string
}
// BaseScreen provides common functionality for installer form screens
type BaseScreen struct {
// Dependencies
controller controller.Controller
styles styles.Styles
window window.Window
// State
initialized bool
hasChanges bool
focusedIndex int
showValues bool
// Form data
fields []FormField
fieldHeights []int
bottomHeight int
// UI components
viewportForm viewport.Model
viewportHelp viewport.Model
formContent string
// Handlers - must be set by concrete implementations
handler BaseScreenHandler
listHandler BaseListHandler // optional, can be nil
// Common utilities
listHelper BaseListHelper
}
// NewBaseScreen creates a new base screen instance
func NewBaseScreen(
c controller.Controller, s styles.Styles, w window.Window,
h BaseScreenHandler, lh BaseListHandler, // can be nil
) *BaseScreen {
return &BaseScreen{
controller: c,
styles: s,
window: w,
showValues: false,
viewportForm: viewport.New(w.GetContentSize()),
viewportHelp: viewport.New(w.GetContentSize()),
handler: h,
listHandler: lh,
fieldHeights: []int{},
listHelper: BaseListHelper{},
}
}
// Init initializes the base screen
func (b *BaseScreen) Init() tea.Cmd {
cmd := b.handler.BuildForm()
b.fields = b.handler.GetFormFields()
b.updateViewports()
return cmd
}
// Update handles common update logic and returns commands only
// Concrete implementations should call this and return themselves as the model
func (b *BaseScreen) Update(msg tea.Msg) tea.Cmd {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
b.updateViewports()
case tea.KeyMsg:
return b.handleKeyPress(msg)
}
return nil
}
// GetFormHotKeys returns the hotkeys for the form
func (b *BaseScreen) GetFormHotKeys() []string {
haveMaskedFields := false
for _, field := range b.fields {
if field.Masked {
haveMaskedFields = true
break
}
}
haveFieldsWithSuggestions := false
for _, field := range b.fields {
if len(field.Suggestions) > 0 {
haveFieldsWithSuggestions = true
break
}
}
hasList := b.listHandler != nil && b.listHandler.GetList() != nil
var hotkeys []string
if len(b.fields) > 0 || hasList {
hotkeys = append(hotkeys, "down|up")
if hasList {
hotkeys = append(hotkeys, "left|right")
}
hotkeys = append(hotkeys, "ctrl+s")
hotkeys = append(hotkeys, "ctrl+r")
}
if haveMaskedFields {
hotkeys = append(hotkeys, "ctrl+h")
}
if haveFieldsWithSuggestions {
hotkeys = append(hotkeys, "tab")
}
hotkeys = append(hotkeys, "enter")
return hotkeys
}
// handleKeyPress handles common keyboard interactions
func (b *BaseScreen) handleKeyPress(msg tea.KeyMsg) tea.Cmd {
switch msg.String() {
case "down":
b.focusNext()
b.updateViewports()
b.ensureFocusVisible()
case "up":
b.focusPrev()
b.updateViewports()
b.ensureFocusVisible()
case "ctrl+s":
return b.saveConfiguration()
case "ctrl+r":
b.resetForm()
b.updateViewports()
case "ctrl+h":
b.toggleShowValues()
b.updateViewports()
case "tab":
b.handleTabCompletion()
case "enter":
return b.saveAndReturn()
}
return nil
}
// View renders the screen
func (b *BaseScreen) View() string {
contentWidth, contentHeight := b.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 {
return locale.UILoading
}
if !b.initialized {
b.handler.BuildForm()
b.fields = b.handler.GetFormFields()
b.updateViewports()
b.initialized = true
}
leftPanel := b.renderForm()
rightPanel := b.renderHelp()
if b.isVerticalLayout() {
return b.renderVerticalLayout(leftPanel, rightPanel, contentWidth, contentHeight)
}
return b.renderHorizontalLayout(leftPanel, rightPanel, contentWidth, contentHeight)
}
// Common form methods
// GetInputWidth calculates the appropriate input width
func (b *BaseScreen) GetInputWidth() int {
viewportWidth, _ := b.getViewportFormSize()
inputWidth := viewportWidth - 6
if b.isVerticalLayout() {
inputWidth = viewportWidth - 4
}
return inputWidth
}
// getViewportFormSize calculates viewport left panel dimensions
func (b *BaseScreen) getViewportFormSize() (int, int) {
contentWidth, contentHeight := b.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 {
return 0, 0
}
if b.isVerticalLayout() {
return contentWidth - PaddingWidth/2, contentHeight - PaddingHeight
} else {
leftWidth, rightWidth := MinMenuWidth, MinInfoWidth
extraWidth := contentWidth - leftWidth - rightWidth - PaddingWidth
if extraWidth > 0 {
leftWidth = min(leftWidth+extraWidth/2, MaxMenuWidth)
}
return leftWidth, contentHeight - PaddingHeight
}
}
// updateViewports updates the viewports with current content
func (b *BaseScreen) updateViewports() {
contentWidth, contentHeight := b.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 {
return
}
b.updateFormContent()
viewportWidth, viewportHeight := b.getViewportFormSize()
formContentHeight := lipgloss.Height(b.formContent)
b.viewportForm.Width = viewportWidth
b.viewportForm.Height = min(viewportHeight, formContentHeight)
b.viewportForm.SetContent(b.formContent) // force update of the viewport content
helpContent := b.renderHelpContent()
b.viewportHelp.Width = lipgloss.Width(helpContent)
b.viewportHelp.Height = lipgloss.Height(helpContent)
b.viewportHelp.SetContent(helpContent) // force update of the viewport content
}
// updateFormContent renders form content and calculates field heights
func (b *BaseScreen) updateFormContent() {
var sections []string
b.fieldHeights = []int{}
inputWidth := b.GetInputWidth()
if b.listHandler != nil {
if listModel := b.listHandler.GetList(); listModel != nil {
listStyle := b.styles.FormInput.Width(inputWidth)
if b.focusedIndex == 0 {
listStyle = listStyle.BorderForeground(styles.Primary)
}
listModel.SetWidth(inputWidth - 4)
renderedList := listStyle.Render(listModel.View())
// field title
titleStyle := b.styles.FormLabel
if b.getFieldIndex() == -1 {
titleStyle = titleStyle.Foreground(styles.Primary)
}
title := titleStyle.Render(b.listHandler.GetListTitle())
sections = append(sections, title)
// field description
description := b.styles.FormHelp.Render(b.listHandler.GetListDescription())
sections = append(sections, description)
sections = append(sections, renderedList)
sections = append(sections, "")
listHeight := lipgloss.Height(b.renderFormContent(sections[:3]))
b.fieldHeights = append(b.fieldHeights, listHeight)
}
}
for i, field := range b.fields {
// check if this field is focused
focused := b.getFieldIndex() == i
// field title
titleStyle := b.styles.FormLabel
if focused {
titleStyle = titleStyle.Foreground(styles.Primary)
}
title := titleStyle.Render(field.Title)
sections = append(sections, title)
// field description
description := b.styles.FormHelp.Render(field.Description)
sections = append(sections, description)
// input field
inputStyle := b.styles.FormInput.Width(inputWidth)
if focused {
inputStyle = inputStyle.BorderForeground(styles.Primary)
}
// configure input
input := field.Input
input.Width = inputWidth - 3
input.SetValue(input.Value()) // force update of the input value
// set up suggestions for tab completion
if len(field.Suggestions) > 0 {
input.ShowSuggestions = true
input.SetSuggestions(field.Suggestions)
}
// apply masking if needed and not showing values
if field.Masked && !b.showValues {
input.EchoMode = textinput.EchoPassword
} else {
input.EchoMode = textinput.EchoNormal
}
// ensure focus state is correct
if focused {
input.Focus()
} else {
input.Blur()
}
renderedInput := inputStyle.Render(input.View())
sections = append(sections, renderedInput)
sections = append(sections, "")
// update the field with configured input
b.fields[i].Input = input
// calculate field height
renderedField := b.renderFormContent([]string{title, description, renderedInput})
b.fieldHeights = append(b.fieldHeights, lipgloss.Height(renderedField))
}
// update list styles
if b.listHandler != nil {
if listModel := b.listHandler.GetList(); listModel != nil {
listModel.Styles.PaginationStyle = b.styles.FormPagination.Width(inputWidth)
}
if listDelegate := b.listHandler.GetListDelegate(); listDelegate != nil {
listDelegate.SetWidth(inputWidth)
}
}
statusMessage := ""
if b.hasChanges {
statusMessage = b.styles.Warning.Render(locale.UIUnsavedChanges)
} else {
statusMessage = b.styles.Success.Render(locale.UIConfigSaved)
}
sections = append(sections, statusMessage)
bottomSections := []string{statusMessage}
if summary := b.handler.GetFormSummary(); summary != "" {
sections = append(sections, "", summary)
bottomSections = append(bottomSections, "", summary)
}
b.bottomHeight = lipgloss.Height(b.renderFormContent(bottomSections))
// update form content
b.formContent = b.renderFormContent(sections)
}
func (b *BaseScreen) renderFormContent(sections []string) string {
content := strings.Join(sections, "\n")
contentWidth, contentHeight := b.window.GetContentSize()
viewportHeight := contentHeight - PaddingHeight // for final rendering
approximateMaxHeight := lipgloss.Height(content)
viewport := viewport.New(contentWidth, max(viewportHeight, approximateMaxHeight*3))
if b.isVerticalLayout() {
xAxisPadding := verticalLayoutPaddings[1] + verticalLayoutPaddings[3]
verticalStyle := lipgloss.NewStyle().Width(contentWidth - xAxisPadding)
content = verticalStyle.Render(content)
} else {
leftWidth, rightWidth := MinMenuWidth, MinInfoWidth
extraWidth := contentWidth - leftWidth - rightWidth - PaddingWidth
if extraWidth > 0 {
leftWidth = min(leftWidth+extraWidth/2, MaxMenuWidth)
}
xAxisPadding := horizontalLayoutPaddings[1] + horizontalLayoutPaddings[3]
content = lipgloss.NewStyle().Width(leftWidth - xAxisPadding).Render(content)
}
viewport.SetContent(content)
viewport.Height = viewport.VisibleLineCount()
return viewport.View()
}
func (b *BaseScreen) renderHelpContent() string {
helpContent := b.handler.GetHelpContent()
if b.isVerticalLayout() {
return b.renderFormContent([]string{helpContent})
}
contentWidth, _ := b.window.GetContentSize()
leftWidth, rightWidth := MinMenuWidth, MinInfoWidth
extraWidth := contentWidth - leftWidth - rightWidth - PaddingWidth
if extraWidth > 0 {
leftWidth = min(leftWidth+extraWidth/2, MaxMenuWidth)
rightWidth = contentWidth - leftWidth - PaddingWidth/2
}
return lipgloss.NewStyle().Width(rightWidth - 2).Render(helpContent)
}
// renderForm renders the left panel with the form
func (b *BaseScreen) renderForm() string {
if !b.initialized {
return locale.UILoading
}
return b.viewportForm.View()
}
// renderHelp renders the right panel with help content
func (b *BaseScreen) renderHelp() string {
if !b.initialized {
return ""
}
return b.viewportHelp.View()
}
// ensureFocusVisible scrolls the viewport to ensure focused field is visible
func (b *BaseScreen) ensureFocusVisible() {
if b.focusedIndex >= len(b.fieldHeights) {
return
}
// calculate y position of focused field
focusY := 0
if b.focusedIndex == len(b.fieldHeights)-1 {
focusY = b.bottomHeight
}
for i := range b.focusedIndex {
focusY += b.fieldHeights[i] + 1 // empty line between fields
}
// get viewport dimensions
visibleRows := b.viewportForm.Height
offset := b.viewportForm.YOffset
// if focused field is above visible area, scroll up
if focusY < offset {
b.viewportForm.YOffset = focusY
}
// if focused field is below visible area, scroll down
if focusY+b.fieldHeights[b.focusedIndex] >= offset+visibleRows {
b.viewportForm.YOffset = focusY + b.fieldHeights[b.focusedIndex] - visibleRows + 1
}
}
// Navigation methods
// focusNext moves focus to the next field
func (b *BaseScreen) focusNext() {
totalElements := b.getTotalElements()
if totalElements == 0 {
return
}
// blur current field
b.blurCurrentField()
// move to next element (with wrapping)
b.focusedIndex = (b.focusedIndex + 1) % totalElements
// focus new field
b.focusCurrentField()
b.updateFormContent()
}
// focusPrev moves focus to the previous field
func (b *BaseScreen) focusPrev() {
totalElements := b.getTotalElements()
if totalElements == 0 {
return
}
// blur current field
b.blurCurrentField()
// move to previous element (with wrapping)
b.focusedIndex = (b.focusedIndex - 1 + totalElements) % totalElements
// focus new field
b.focusCurrentField()
b.updateFormContent()
}
// getTotalElements returns the total number of navigable elements
func (b *BaseScreen) getTotalElements() int {
total := len(b.fields)
// add 1 for list if present
if b.listHandler != nil && b.listHandler.GetList() != nil {
total++
}
return total
}
// blurCurrentField removes focus from the currently focused field
func (b *BaseScreen) blurCurrentField() {
fieldIndex := b.getFieldIndex()
if fieldIndex >= 0 && fieldIndex < len(b.fields) {
b.fields[fieldIndex].Input.Blur()
}
}
// focusCurrentField sets focus on the currently focused field
func (b *BaseScreen) focusCurrentField() {
fieldIndex := b.getFieldIndex()
if fieldIndex >= 0 && fieldIndex < len(b.fields) {
b.fields[fieldIndex].Input.Focus()
}
}
// getFieldIndex returns the field index for the current focusedIndex (-1 if focused on list)
func (b *BaseScreen) getFieldIndex() int {
if b.listHandler != nil && b.listHandler.GetList() != nil {
// list is at index 0, fields start at index 1
return b.focusedIndex - 1
}
// no list, fields start at index 0
return b.focusedIndex
}
// toggleShowValues toggles visibility of masked values
func (b *BaseScreen) toggleShowValues() {
b.showValues = !b.showValues
b.updateFormContent()
}
// handleTabCompletion handles tab completion for focused field
func (b *BaseScreen) handleTabCompletion() {
fieldIndex := b.getFieldIndex()
// check if we're focused on a valid field
if fieldIndex >= 0 && fieldIndex < len(b.fields) {
field := &b.fields[fieldIndex]
// only handle tab completion if field has suggestions
if len(field.Suggestions) > 0 {
// use textinput's built-in suggestion functionality
if suggestion := field.Input.CurrentSuggestion(); suggestion != "" {
oldValue := field.Input.Value()
field.Input.SetValue(suggestion)
field.Input.CursorEnd()
field.Value = suggestion
b.hasChanges = true
// notify handler about the change
b.handler.OnFieldChanged(fieldIndex, oldValue, suggestion)
// update the fields array
b.fields[fieldIndex] = *field
b.handler.SetFormFields(b.fields)
b.updateViewports()
}
}
}
}
// resetForm resets the form to default values
func (b *BaseScreen) resetForm() {
b.handler.HandleReset()
b.fields = b.handler.GetFormFields()
b.hasChanges = false
b.updateFormContent()
}
// saveConfiguration saves the current configuration
func (b *BaseScreen) saveConfiguration() tea.Cmd {
if err := b.handler.HandleSave(); err != nil {
logger.Errorf("[BaseScreen] SAVE: error: %v", err)
return nil
}
b.hasChanges = false
b.updateViewports()
return nil
}
// saveAndReturn saves and returns to previous screen
func (b *BaseScreen) saveAndReturn() tea.Cmd {
// save first
cmd := b.saveConfiguration()
if cmd != nil {
return cmd
}
// return to previous screen
return func() tea.Msg {
return NavigationMsg{GoBack: true}
}
}
// Layout methods
// isVerticalLayout determines if vertical layout should be used
func (b *BaseScreen) isVerticalLayout() bool {
contentWidth := b.window.GetContentWidth()
return contentWidth < (MinMenuWidth + MinInfoWidth + PaddingWidth)
}
// renderVerticalLayout renders content in vertical layout
func (b *BaseScreen) renderVerticalLayout(leftPanel, rightPanel string, width, height int) string {
verticalStyle := lipgloss.NewStyle().Width(width).Padding(verticalLayoutPaddings...)
leftStyled := verticalStyle.Render(leftPanel)
rightStyled := verticalStyle.Render(rightPanel)
if lipgloss.Height(leftStyled)+lipgloss.Height(rightStyled)+3 < height {
return lipgloss.JoinVertical(lipgloss.Left,
verticalStyle.Render(leftPanel),
verticalStyle.Height(2).Render("\n"),
verticalStyle.Render(rightPanel),
)
}
return verticalStyle.Render(leftPanel)
}
// renderHorizontalLayout renders content in horizontal layout
func (b *BaseScreen) renderHorizontalLayout(leftPanel, rightPanel string, width, height int) string {
leftWidth, rightWidth := MinMenuWidth, MinInfoWidth
extraWidth := width - leftWidth - rightWidth - PaddingWidth
if extraWidth > 0 {
leftWidth = min(leftWidth+extraWidth/2, MaxMenuWidth)
rightWidth = width - leftWidth - PaddingWidth/2
}
leftStyled := lipgloss.NewStyle().Width(leftWidth).Padding(horizontalLayoutPaddings...).Render(leftPanel)
rightStyled := lipgloss.NewStyle().Width(rightWidth).PaddingLeft(2).Render(rightPanel)
viewport := viewport.New(width, height-PaddingHeight)
viewport.SetContent(lipgloss.JoinHorizontal(lipgloss.Top, leftStyled, rightStyled))
return viewport.View()
}
// Helper methods for concrete implementations
// HandleFieldInput handles input for a specific field
func (b *BaseScreen) HandleFieldInput(msg tea.KeyMsg) tea.Cmd {
fieldIndex := b.getFieldIndex()
// all hotkeys are handled by handleKeyPress(msg tea.KeyMsg) method
// inherit screen must call HandleUpdate(msg tea.Msg) for all uncaught messages
if slices.Contains(b.GetFormHotKeys(), msg.String()) {
return nil
}
// check if we're focused on a valid field
if fieldIndex >= 0 && fieldIndex < len(b.fields) {
var cmd tea.Cmd
oldValue := b.fields[fieldIndex].Input.Value()
b.fields[fieldIndex].Input, cmd = b.fields[fieldIndex].Input.Update(msg)
newValue := b.fields[fieldIndex].Input.Value()
if oldValue != newValue {
b.fields[fieldIndex].Value = newValue
b.hasChanges = true
b.handler.OnFieldChanged(fieldIndex, oldValue, newValue)
}
b.updateViewports()
return cmd
}
return nil
}
// HandleListInput handles input for the list component
func (b *BaseScreen) HandleListInput(msg tea.KeyMsg) tea.Cmd {
// check if we have a list and we're focused on it (skip if not)
if b.listHandler == nil {
return nil
}
// check if focused on list (index 0 when list is present)
isFocusedOnList := b.listHandler.GetList() != nil && b.focusedIndex == 0
if !isFocusedOnList {
return nil
}
// filter list input keys to slide the list
switch msg.String() {
case "left", "right":
break
default:
return nil
}
listModel := b.listHandler.GetList()
if listModel == nil {
return nil
}
// get old selection
oldSelection := ""
if selectedItem := listModel.SelectedItem(); selectedItem != nil {
oldSelection = selectedItem.FilterValue()
}
// update list
var cmd tea.Cmd
*listModel, cmd = listModel.Update(msg)
// get new selection
newSelection := ""
if selectedItem := listModel.SelectedItem(); selectedItem != nil {
newSelection = selectedItem.FilterValue()
}
// notify handler if selection changed
if oldSelection != newSelection {
b.listHandler.OnListSelectionChanged(oldSelection, newSelection)
b.hasChanges = true
b.updateViewports()
}
return cmd
}
// GetController returns the state controller
func (b *BaseScreen) GetController() controller.Controller {
return b.controller
}
// GetStyles returns the styles
func (b *BaseScreen) GetStyles() styles.Styles {
return b.styles
}
// GetWindow returns the window
func (b *BaseScreen) GetWindow() window.Window {
return b.window
}
// SetHasChanges sets the hasChanges flag
func (b *BaseScreen) SetHasChanges(hasChanges bool) {
b.hasChanges = hasChanges
}
// GetHasChanges returns the hasChanges flag
func (b *BaseScreen) GetHasChanges() bool {
return b.hasChanges
}
// GetShowValues returns the showValues flag
func (b *BaseScreen) GetShowValues() bool {
return b.showValues
}
// GetFocusedIndex returns the currently focused field index
func (b *BaseScreen) GetFocusedIndex() int {
return b.focusedIndex
}
// SetFocusedIndex sets the focused field index
func (b *BaseScreen) SetFocusedIndex(index int) {
b.focusedIndex = index
}
// GetListHelper returns the list helper utility
func (b *BaseScreen) GetListHelper() *BaseListHelper {
return &b.listHelper
}
backend/cmd/installer/wizard/models/docker_form.go
+445
View File
@@ -0,0 +1,445 @@
package models
import (
"fmt"
"os"
"strings"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
tea "github.com/charmbracelet/bubbletea"
)
// DockerFormModel represents the Docker Environment configuration form
type DockerFormModel struct {
*BaseScreen
}
// NewDockerFormModel creates a new Docker Environment form model
func NewDockerFormModel(c controller.Controller, s styles.Styles, w window.Window) *DockerFormModel {
m := &DockerFormModel{}
// create base screen with this model as handler (no list handler needed)
m.BaseScreen = NewBaseScreen(c, s, w, m, nil)
return m
}
// BaseScreenHandler interface implementation
func (m *DockerFormModel) BuildForm() tea.Cmd {
config := m.GetController().GetDockerConfig()
fields := []FormField{}
// Container capabilities
fields = append(fields, m.createBooleanField("docker_inside",
locale.ToolsDockerInside,
locale.ToolsDockerInsideDesc,
config.DockerInside,
))
fields = append(fields, m.createBooleanField("docker_net_admin",
locale.ToolsDockerNetAdmin,
locale.ToolsDockerNetAdminDesc,
config.DockerNetAdmin,
))
// Connection settings
fields = append(fields, m.createTextField("docker_socket",
locale.ToolsDockerSocket,
locale.ToolsDockerSocketDesc,
config.DockerSocket,
false,
))
fields = append(fields, m.createTextField("docker_network",
locale.ToolsDockerNetwork,
locale.ToolsDockerNetworkDesc,
config.DockerNetwork,
false,
))
fields = append(fields, m.createTextField("docker_public_ip",
locale.ToolsDockerPublicIP,
locale.ToolsDockerPublicIPDesc,
config.DockerPublicIP,
false,
))
// Storage configuration
fields = append(fields, m.createTextField("docker_work_dir",
locale.ToolsDockerWorkDir,
locale.ToolsDockerWorkDirDesc,
config.DockerWorkDir,
false,
))
// Default images
fields = append(fields, m.createTextField("docker_default_image",
locale.ToolsDockerDefaultImage,
locale.ToolsDockerDefaultImageDesc,
config.DockerDefaultImage,
false,
))
fields = append(fields, m.createTextField("docker_default_image_for_pentest",
locale.ToolsDockerDefaultImageForPentest,
locale.ToolsDockerDefaultImageForPentestDesc,
config.DockerDefaultImageForPentest,
false,
))
// TLS connection settings (optional)
fields = append(fields, m.createTextField("docker_host",
locale.ToolsDockerHost,
locale.ToolsDockerHostDesc,
config.DockerHost,
false,
))
fields = append(fields, m.createBooleanField("docker_tls_verify",
locale.ToolsDockerTLSVerify,
locale.ToolsDockerTLSVerifyDesc,
config.DockerTLSVerify,
))
fields = append(fields, m.createTextField("docker_cert_path",
locale.ToolsDockerCertPath,
locale.ToolsDockerCertPathDesc,
config.HostDockerCertPath,
false,
))
m.SetFormFields(fields)
return nil
}
func (m *DockerFormModel) createBooleanField(key, title, description string, envVar loader.EnvVar) FormField {
input := NewBooleanInput(m.GetStyles(), m.GetWindow(), envVar)
return FormField{
Key: key,
Title: title,
Description: description,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
Suggestions: input.AvailableSuggestions(),
}
}
func (m *DockerFormModel) createTextField(key, title, description string, envVar loader.EnvVar, masked bool) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), envVar)
return FormField{
Key: key,
Title: title,
Description: description,
Required: false,
Masked: masked,
Input: input,
Value: input.Value(),
}
}
func (m *DockerFormModel) GetFormTitle() string {
return locale.ToolsDockerFormTitle
}
func (m *DockerFormModel) GetFormDescription() string {
return locale.ToolsDockerFormDescription
}
func (m *DockerFormModel) GetFormName() string {
return locale.ToolsDockerFormName
}
func (m *DockerFormModel) GetFormSummary() string {
return ""
}
func (m *DockerFormModel) GetFormOverview() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.ToolsDockerFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.ToolsDockerFormDescription))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Render(locale.ToolsDockerFormOverview))
return strings.Join(sections, "\n")
}
func (m *DockerFormModel) GetCurrentConfiguration() string {
var sections []string
config := m.GetController().GetDockerConfig()
sections = append(sections, m.GetStyles().Subtitle.Render(m.GetFormName()))
// Container capabilities
dockerInside := config.DockerInside.Value
if dockerInside == "" {
dockerInside = config.DockerInside.Default
}
if dockerInside == "true" {
sections = append(sections, fmt.Sprintf("• Docker Access: %s",
m.GetStyles().Success.Render(locale.StatusEnabled)))
} else {
sections = append(sections, fmt.Sprintf("• Docker Access: %s",
m.GetStyles().Warning.Render(locale.StatusDisabled)))
}
dockerNetAdmin := config.DockerNetAdmin.Value
if dockerNetAdmin == "" {
dockerNetAdmin = config.DockerNetAdmin.Default
}
if dockerNetAdmin == "true" {
sections = append(sections, fmt.Sprintf("• Network Admin: %s",
m.GetStyles().Success.Render(locale.StatusEnabled)))
} else {
sections = append(sections, fmt.Sprintf("• Network Admin: %s",
m.GetStyles().Warning.Render(locale.StatusDisabled)))
}
// Connection settings
if config.DockerNetwork.Value != "" {
sections = append(sections, fmt.Sprintf("• Custom Network: %s",
m.GetStyles().Success.Render(locale.StatusConfigured)))
} else {
sections = append(sections, fmt.Sprintf("• Custom Network: %s",
m.GetStyles().Warning.Render(locale.StatusNotConfigured)))
}
if config.DockerPublicIP.Value != "" && config.DockerPublicIP.Value != "0.0.0.0" {
sections = append(sections, fmt.Sprintf("• Public IP: %s",
m.GetStyles().Success.Render(locale.StatusConfigured)))
} else {
sections = append(sections, fmt.Sprintf("• Public IP: %s",
m.GetStyles().Warning.Render(locale.StatusNotConfigured)))
}
// Default images
if config.DockerDefaultImage.Value != "" {
sections = append(sections, fmt.Sprintf("• Default Image: %s",
m.GetStyles().Success.Render(locale.StatusConfigured)))
} else {
sections = append(sections, fmt.Sprintf("• Default Image: %s",
m.GetStyles().Warning.Render(locale.StatusNotConfigured)))
}
if config.DockerDefaultImageForPentest.Value != "" {
sections = append(sections, fmt.Sprintf("• Pentest Image: %s",
m.GetStyles().Success.Render(locale.StatusConfigured)))
} else {
sections = append(sections, fmt.Sprintf("• Pentest Image: %s",
m.GetStyles().Warning.Render(locale.StatusNotConfigured)))
}
// TLS settings
if config.DockerHost.Value != "" && config.DockerTLSVerify.Value == "1" && config.HostDockerCertPath.Value != "" {
sections = append(sections, fmt.Sprintf("• TLS Connection: %s",
m.GetStyles().Success.Render(locale.StatusConfigured)))
} else if config.DockerHost.Value != "" {
sections = append(sections, fmt.Sprintf("• Remote Connection: %s",
m.GetStyles().Success.Render(locale.StatusConfigured)))
}
sections = append(sections, "")
if config.Configured {
sections = append(sections, m.GetStyles().Success.Render(locale.MessageDockerConfigured))
} else {
sections = append(sections, m.GetStyles().Warning.Render(locale.MessageDockerNotConfigured))
}
return strings.Join(sections, "\n")
}
func (m *DockerFormModel) IsConfigured() bool {
return m.GetController().GetDockerConfig().Configured
}
func (m *DockerFormModel) GetHelpContent() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.ToolsDockerFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.ToolsDockerFormDescription))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Render(locale.ToolsDockerGeneralHelp))
sections = append(sections, "")
// Show field-specific help based on focused field
fieldIndex := m.GetFocusedIndex()
fields := m.GetFormFields()
if fieldIndex >= 0 && fieldIndex < len(fields) {
field := fields[fieldIndex]
switch field.Key {
case "docker_inside":
sections = append(sections, locale.ToolsDockerInsideHelp)
case "docker_net_admin":
sections = append(sections, locale.ToolsDockerNetAdminHelp)
case "docker_socket":
sections = append(sections, locale.ToolsDockerSocketHelp)
case "docker_network":
sections = append(sections, locale.ToolsDockerNetworkHelp)
case "docker_public_ip":
sections = append(sections, locale.ToolsDockerPublicIPHelp)
case "docker_work_dir":
sections = append(sections, locale.ToolsDockerWorkDirHelp)
case "docker_default_image":
sections = append(sections, locale.ToolsDockerDefaultImageHelp)
case "docker_default_image_for_pentest":
sections = append(sections, locale.ToolsDockerDefaultImageForPentestHelp)
case "docker_host":
sections = append(sections, locale.ToolsDockerHostHelp)
case "docker_tls_verify":
sections = append(sections, locale.ToolsDockerTLSVerifyHelp)
case "docker_cert_path":
sections = append(sections, locale.ToolsDockerCertPathHelp)
default:
sections = append(sections, locale.ToolsDockerFormOverview)
}
}
return strings.Join(sections, "\n")
}
func (m *DockerFormModel) HandleSave() error {
config := m.GetController().GetDockerConfig()
fields := m.GetFormFields()
// create a working copy of the current config to modify
newConfig := &controller.DockerConfig{
// copy current EnvVar fields - they preserve metadata like Line, IsPresent, etc.
DockerInside: config.DockerInside,
DockerNetAdmin: config.DockerNetAdmin,
DockerSocket: config.DockerSocket,
DockerNetwork: config.DockerNetwork,
DockerPublicIP: config.DockerPublicIP,
DockerWorkDir: config.DockerWorkDir,
DockerDefaultImage: config.DockerDefaultImage,
DockerDefaultImageForPentest: config.DockerDefaultImageForPentest,
DockerHost: config.DockerHost,
DockerTLSVerify: config.DockerTLSVerify,
HostDockerCertPath: config.HostDockerCertPath,
}
// update field values based on form input
for _, field := range fields {
value := strings.TrimSpace(field.Input.Value())
switch field.Key {
case "docker_inside":
// validate boolean input
if value != "" && value != "true" && value != "false" {
return fmt.Errorf("invalid boolean value for Docker Access: %s (must be 'true' or 'false')", value)
}
newConfig.DockerInside.Value = value
case "docker_net_admin":
// validate boolean input
if value != "" && value != "true" && value != "false" {
return fmt.Errorf("invalid boolean value for Network Admin: %s (must be 'true' or 'false')", value)
}
newConfig.DockerNetAdmin.Value = value
case "docker_socket":
newConfig.DockerSocket.Value = value
case "docker_network":
newConfig.DockerNetwork.Value = value
case "docker_public_ip":
newConfig.DockerPublicIP.Value = value
case "docker_work_dir":
newConfig.DockerWorkDir.Value = value
case "docker_default_image":
newConfig.DockerDefaultImage.Value = value
case "docker_default_image_for_pentest":
newConfig.DockerDefaultImageForPentest.Value = value
case "docker_host":
newConfig.DockerHost.Value = value
case "docker_tls_verify":
// validate boolean input for TLS verification
if value != "" && value != "true" && value != "false" && value != "1" && value != "0" {
return fmt.Errorf("invalid boolean value for TLS Verification: %s (must be 'true', 'false', '1', or '0')", value)
}
// normalize to "1" or "" for TLS verification
switch value {
case "true":
value = "1"
case "false":
value = ""
}
newConfig.DockerTLSVerify.Value = value
case "docker_cert_path":
// validate cert path if provided
if value != "" {
info, err := os.Stat(value)
if err != nil {
if os.IsNotExist(err) {
return fmt.Errorf("docker cert path does not exist: %s", value)
}
return fmt.Errorf("cannot access docker cert path %s: %v", value, err)
}
if !info.IsDir() {
return fmt.Errorf("docker cert path must be a directory, not a file: %s", value)
}
}
newConfig.HostDockerCertPath.Value = value
}
}
// save the configuration
if err := m.GetController().UpdateDockerConfig(newConfig); err != nil {
logger.Errorf("[DockerFormModel] SAVE: error updating Docker config: %v", err)
return err
}
logger.Log("[DockerFormModel] SAVE: success")
return nil
}
func (m *DockerFormModel) HandleReset() {
// reset config to defaults
m.GetController().ResetDockerConfig()
// rebuild form with reset values
m.BuildForm()
}
func (m *DockerFormModel) OnFieldChanged(fieldIndex int, oldValue, newValue string) {
// additional validation could be added here if needed
}
func (m *DockerFormModel) GetFormFields() []FormField {
return m.BaseScreen.fields
}
func (m *DockerFormModel) SetFormFields(fields []FormField) {
m.BaseScreen.fields = fields
}
// Update method - handle screen-specific input
func (m *DockerFormModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.KeyMsg:
// then handle field input
if cmd := m.HandleFieldInput(msg); cmd != nil {
return m, cmd
}
}
// delegate to base screen for common handling
cmd := m.BaseScreen.Update(msg)
return m, cmd
}
// Compile-time interface validation
var _ BaseScreenModel = (*DockerFormModel)(nil)
var _ BaseScreenHandler = (*DockerFormModel)(nil)
backend/cmd/installer/wizard/models/embedder_form.go
+637
View File
@@ -0,0 +1,637 @@
package models
import (
"fmt"
"strconv"
"strings"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/list"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
// EmbeddingProviderInfo contains information about an embedding provider
type EmbeddingProviderInfo struct {
ID string
Name string
Description string
URLPlaceholder string
APIKeyPlaceholder string
ModelPlaceholder string
RequiresAPIKey bool
SupportsURL bool
SupportsModel bool
HelpText string
}
// EmbedderFormModel represents the Embedder configuration form
type EmbedderFormModel struct {
*BaseScreen
// screen-specific components
providerList list.Model
providerDelegate *BaseListDelegate
// provider information
providers map[string]*EmbeddingProviderInfo
}
// NewEmbedderFormModel creates a new Embedder form model
func NewEmbedderFormModel(c controller.Controller, s styles.Styles, w window.Window) *EmbedderFormModel {
m := &EmbedderFormModel{
providers: initEmbeddingProviders(),
}
m.BaseScreen = NewBaseScreen(c, s, w, m, m)
m.initializeProviderList(s)
return m
}
// initEmbeddingProviders initializes the provider information
func initEmbeddingProviders() map[string]*EmbeddingProviderInfo {
return map[string]*EmbeddingProviderInfo{
locale.EmbedderProviderIDDefault: {
ID: locale.EmbedderProviderIDDefault,
Name: locale.EmbedderProviderDefault,
Description: locale.EmbedderProviderDefaultDesc,
URLPlaceholder: "",
APIKeyPlaceholder: "",
ModelPlaceholder: "",
RequiresAPIKey: false,
SupportsURL: false,
SupportsModel: false,
HelpText: locale.EmbedderHelpDefault,
},
locale.EmbedderProviderIDOpenAI: {
ID: locale.EmbedderProviderIDOpenAI,
Name: locale.EmbedderProviderOpenAI,
Description: locale.EmbedderProviderOpenAIDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderOpenAI,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderDefault,
ModelPlaceholder: locale.EmbedderModelPlaceholderOpenAI,
RequiresAPIKey: true,
SupportsURL: true,
SupportsModel: true,
HelpText: locale.EmbedderHelpOpenAI,
},
locale.EmbedderProviderIDOllama: {
ID: locale.EmbedderProviderIDOllama,
Name: locale.EmbedderProviderOllama,
Description: locale.EmbedderProviderOllamaDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderOllama,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderOllama,
ModelPlaceholder: locale.EmbedderModelPlaceholderOllama,
RequiresAPIKey: false,
SupportsURL: true,
SupportsModel: true,
HelpText: locale.EmbedderHelpOllama,
},
locale.EmbedderProviderIDMistral: {
ID: locale.EmbedderProviderIDMistral,
Name: locale.EmbedderProviderMistral,
Description: locale.EmbedderProviderMistralDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderMistral,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderMistral,
ModelPlaceholder: locale.EmbedderModelPlaceholderMistral,
RequiresAPIKey: true,
SupportsURL: true,
SupportsModel: false,
HelpText: locale.EmbedderHelpMistral,
},
locale.EmbedderProviderIDJina: {
ID: locale.EmbedderProviderIDJina,
Name: locale.EmbedderProviderJina,
Description: locale.EmbedderProviderJinaDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderJina,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderJina,
ModelPlaceholder: locale.EmbedderModelPlaceholderJina,
RequiresAPIKey: true,
SupportsURL: true,
SupportsModel: true,
HelpText: locale.EmbedderHelpJina,
},
locale.EmbedderProviderIDHuggingFace: {
ID: locale.EmbedderProviderIDHuggingFace,
Name: locale.EmbedderProviderHuggingFace,
Description: locale.EmbedderProviderHuggingFaceDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderHuggingFace,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderHuggingFace,
ModelPlaceholder: locale.EmbedderModelPlaceholderHuggingFace,
RequiresAPIKey: true,
SupportsURL: true,
SupportsModel: true,
HelpText: locale.EmbedderHelpHuggingFace,
},
locale.EmbedderProviderIDGoogleAI: {
ID: locale.EmbedderProviderIDGoogleAI,
Name: locale.EmbedderProviderGoogleAI,
Description: locale.EmbedderProviderGoogleAIDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderGoogleAI,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderGoogleAI,
ModelPlaceholder: locale.EmbedderModelPlaceholderGoogleAI,
RequiresAPIKey: true,
SupportsURL: false,
SupportsModel: true,
HelpText: locale.EmbedderHelpGoogleAI,
},
locale.EmbedderProviderIDVoyageAI: {
ID: locale.EmbedderProviderIDVoyageAI,
Name: locale.EmbedderProviderVoyageAI,
Description: locale.EmbedderProviderVoyageAIDesc,
URLPlaceholder: locale.EmbedderURLPlaceholderVoyageAI,
APIKeyPlaceholder: locale.EmbedderAPIKeyPlaceholderVoyageAI,
ModelPlaceholder: locale.EmbedderModelPlaceholderVoyageAI,
RequiresAPIKey: true,
SupportsURL: false,
SupportsModel: true,
HelpText: locale.EmbedderHelpVoyageAI,
},
locale.EmbedderProviderIDDisabled: {
ID: locale.EmbedderProviderIDDisabled,
Name: locale.EmbedderProviderDisabled,
Description: locale.EmbedderProviderDisabledDesc,
URLPlaceholder: "",
APIKeyPlaceholder: "",
ModelPlaceholder: "",
RequiresAPIKey: false,
SupportsURL: false,
SupportsModel: false,
HelpText: locale.EmbedderHelpDisabled,
},
}
}
// initializeProviderList sets up the provider selection list
func (m *EmbedderFormModel) initializeProviderList(styles styles.Styles) {
options := []BaseListOption{
{Value: locale.EmbedderProviderIDDefault, Display: locale.EmbedderProviderDefault},
{Value: locale.EmbedderProviderIDOpenAI, Display: locale.EmbedderProviderOpenAI},
{Value: locale.EmbedderProviderIDOllama, Display: locale.EmbedderProviderOllama},
{Value: locale.EmbedderProviderIDMistral, Display: locale.EmbedderProviderMistral},
{Value: locale.EmbedderProviderIDJina, Display: locale.EmbedderProviderJina},
{Value: locale.EmbedderProviderIDHuggingFace, Display: locale.EmbedderProviderHuggingFace},
{Value: locale.EmbedderProviderIDGoogleAI, Display: locale.EmbedderProviderGoogleAI},
{Value: locale.EmbedderProviderIDVoyageAI, Display: locale.EmbedderProviderVoyageAI},
{Value: locale.EmbedderProviderIDDisabled, Display: locale.EmbedderProviderDisabled},
}
m.providerDelegate = NewBaseListDelegate(
styles.FormLabel.Align(lipgloss.Center),
MinMenuWidth-6,
)
m.providerList = m.GetListHelper().CreateList(options, m.providerDelegate, MinMenuWidth-6, 3)
// set current selection
config := m.GetController().GetEmbedderConfig()
selectedProvider := m.getProviderID(config.Provider.Value)
m.GetListHelper().SelectByValue(&m.providerList, selectedProvider)
}
// getProviderID converts provider value to ID
func (m *EmbedderFormModel) getProviderID(provider string) string {
switch provider {
case "":
return locale.EmbedderProviderIDDefault
case locale.EmbedderProviderIDDisabled:
return locale.EmbedderProviderIDDisabled
default:
// check if it's a known provider
if _, exists := m.providers[provider]; exists {
return provider
}
// fallback to default for unknown providers
return locale.EmbedderProviderIDDefault
}
}
// getSelectedProvider returns the currently selected provider ID
func (m *EmbedderFormModel) getSelectedProvider() string {
selectedValue := m.GetListHelper().GetSelectedValue(&m.providerList)
if selectedValue == "" {
return locale.EmbedderProviderIDDefault
}
return selectedValue
}
// getCurrentProviderInfo returns information about the currently selected provider
func (m *EmbedderFormModel) getCurrentProviderInfo() *EmbeddingProviderInfo {
providerID := m.getSelectedProvider()
if info, exists := m.providers[providerID]; exists {
return info
}
return m.providers[locale.EmbedderProviderIDDefault]
}
// BaseScreenHandler interface implementation
func (m *EmbedderFormModel) BuildForm() tea.Cmd {
config := m.GetController().GetEmbedderConfig()
fields := []FormField{}
providerInfo := m.getCurrentProviderInfo()
// URL field (if supported)
if providerInfo.SupportsURL {
fields = append(fields, m.createURLField(config, providerInfo))
}
// API Key field (if required)
if providerInfo.RequiresAPIKey {
fields = append(fields, m.createAPIKeyField(config, providerInfo))
}
// Model field (if supported)
if providerInfo.SupportsModel {
fields = append(fields, m.createModelField(config, providerInfo))
}
// Batch size field (always show except for disabled)
if providerInfo.ID != locale.EmbedderProviderIDDisabled {
fields = append(fields, m.createBatchSizeField(config))
}
// Strip newlines field (always show except for disabled)
if providerInfo.ID != locale.EmbedderProviderIDDisabled {
fields = append(fields, m.createStripNewLinesField(config))
}
// Max text bytes field (always show except for disabled)
if providerInfo.ID != locale.EmbedderProviderIDDisabled {
fields = append(fields, m.createMaxTextBytesField(config))
}
m.SetFormFields(fields)
return nil
}
func (m *EmbedderFormModel) createURLField(
config *controller.EmbedderConfig, providerInfo *EmbeddingProviderInfo,
) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), config.URL)
if providerInfo.URLPlaceholder != "" {
input.Placeholder = providerInfo.URLPlaceholder
}
return FormField{
Key: "url",
Title: locale.EmbedderFormURL,
Description: locale.EmbedderFormURLDesc,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
}
}
func (m *EmbedderFormModel) createAPIKeyField(
config *controller.EmbedderConfig, providerInfo *EmbeddingProviderInfo,
) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), config.APIKey)
if providerInfo.APIKeyPlaceholder != "" {
input.Placeholder = providerInfo.APIKeyPlaceholder
}
return FormField{
Key: "api_key",
Title: locale.EmbedderFormAPIKey,
Description: locale.EmbedderFormAPIKeyDesc,
Required: false,
Masked: true,
Input: input,
Value: input.Value(),
}
}
func (m *EmbedderFormModel) createModelField(
config *controller.EmbedderConfig, providerInfo *EmbeddingProviderInfo,
) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), config.Model)
if providerInfo.ModelPlaceholder != "" {
input.Placeholder = providerInfo.ModelPlaceholder
}
return FormField{
Key: "model",
Title: locale.EmbedderFormModel,
Description: locale.EmbedderFormModelDesc,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
}
}
func (m *EmbedderFormModel) createBatchSizeField(config *controller.EmbedderConfig) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), config.BatchSize)
return FormField{
Key: "batch_size",
Title: locale.EmbedderFormBatchSize,
Description: locale.EmbedderFormBatchSizeDesc,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
}
}
func (m *EmbedderFormModel) createStripNewLinesField(config *controller.EmbedderConfig) FormField {
input := NewBooleanInput(m.GetStyles(), m.GetWindow(), config.StripNewLines)
return FormField{
Key: "strip_newlines",
Title: locale.EmbedderFormStripNewLines,
Description: locale.EmbedderFormStripNewLinesDesc,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
Suggestions: input.AvailableSuggestions(),
}
}
func (m *EmbedderFormModel) createMaxTextBytesField(config *controller.EmbedderConfig) FormField {
input := NewTextInput(m.GetStyles(), m.GetWindow(), config.MaxTextBytes)
return FormField{
Key: "max_text_bytes",
Title: locale.EmbedderFormMaxTextBytes,
Description: locale.EmbedderFormMaxTextBytesDesc,
Required: false,
Masked: false,
Input: input,
Value: input.Value(),
}
}
func (m *EmbedderFormModel) GetFormTitle() string {
return locale.EmbedderFormTitle
}
func (m *EmbedderFormModel) GetFormDescription() string {
return locale.EmbedderFormDescription
}
func (m *EmbedderFormModel) GetFormName() string {
return locale.EmbedderFormName
}
func (m *EmbedderFormModel) GetFormSummary() string {
return ""
}
func (m *EmbedderFormModel) GetFormOverview() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.EmbedderFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.EmbedderFormDescription))
sections = append(sections, "")
sections = append(sections, locale.EmbedderFormOverview)
return strings.Join(sections, "\n")
}
func (m *EmbedderFormModel) GetCurrentConfiguration() string {
var sections []string
config := m.GetController().GetEmbedderConfig()
sections = append(sections, m.GetStyles().Subtitle.Render(m.GetFormName()))
providerID := m.getProviderID(config.Provider.Value)
providerInfo := m.providers[providerID]
if config.Configured {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormProvider, m.GetStyles().Success.Render(providerInfo.Name)))
} else {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormProvider, m.GetStyles().Warning.Render(providerInfo.Name+
" ("+locale.StatusNotConfigured+")")))
}
if config.URL.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormURL, m.GetStyles().Info.Render(config.URL.Value)))
}
if config.APIKey.Value != "" {
maskedKey := strings.Repeat("*", len(config.APIKey.Value))
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormAPIKey, m.GetStyles().Muted.Render(maskedKey)))
}
if config.Model.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormModel, m.GetStyles().Info.Render(config.Model.Value)))
}
if config.BatchSize.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormBatchSize, m.GetStyles().Info.Render(config.BatchSize.Value)))
} else if config.BatchSize.Default != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormBatchSize, m.GetStyles().Info.Render(config.BatchSize.Default)))
}
stripNewLines := config.StripNewLines.Value
if stripNewLines == "" {
stripNewLines = config.StripNewLines.Default
}
if stripNewLines != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormStripNewLines, m.GetStyles().Info.Render(stripNewLines)))
}
maxTextBytes := config.MaxTextBytes.Value
if maxTextBytes == "" {
maxTextBytes = config.MaxTextBytes.Default
}
if maxTextBytes != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.EmbedderFormMaxTextBytes, m.GetStyles().Info.Render(maxTextBytes)))
}
return strings.Join(sections, "\n")
}
func (m *EmbedderFormModel) IsConfigured() bool {
return m.GetController().GetEmbedderConfig().Configured
}
func (m *EmbedderFormModel) GetHelpContent() string {
var sections []string
providerInfo := m.getCurrentProviderInfo()
config := m.GetController().GetEmbedderConfig()
sections = append(sections, m.GetStyles().Subtitle.Render(locale.EmbedderFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.EmbedderFormDescription))
sections = append(sections, "")
if config.Configured {
sections = append(sections, fmt.Sprintf("%s %s\n%s",
m.GetStyles().Warning.Bold(true).Render(locale.EmbedderHelpAttentionPrefix),
m.GetStyles().Paragraph.Render(locale.EmbedderHelpAttention),
m.GetStyles().Warning.Render(locale.EmbedderHelpAttentionSuffix)))
sections = append(sections, "")
}
sections = append(sections, m.GetStyles().Paragraph.Render(locale.EmbedderHelpGeneral))
sections = append(sections, "")
sections = append(sections, providerInfo.HelpText)
return strings.Join(sections, "\n")
}
func (m *EmbedderFormModel) HandleSave() error {
config := m.GetController().GetEmbedderConfig()
selectedProvider := m.getSelectedProvider()
fields := m.GetFormFields()
// create a working copy of the current config to modify
newConfig := &controller.EmbedderConfig{
// copy current EnvVar fields - they preserve metadata like Line, IsPresent, etc.
Provider: config.Provider,
URL: config.URL,
APIKey: config.APIKey,
Model: config.Model,
BatchSize: config.BatchSize,
StripNewLines: config.StripNewLines,
MaxTextBytes: config.MaxTextBytes,
}
// set provider
switch selectedProvider {
case locale.EmbedderProviderIDDefault:
newConfig.Provider.Value = "" // empty means use default (openai)
case locale.EmbedderProviderIDDisabled:
newConfig.Provider.Value = locale.EmbedderProviderIDDisabled
default:
newConfig.Provider.Value = selectedProvider
}
// update field values based on form input
for _, field := range fields {
value := strings.TrimSpace(field.Input.Value())
switch field.Key {
case "url":
newConfig.URL.Value = value
case "api_key":
newConfig.APIKey.Value = value
case "model":
newConfig.Model.Value = value
case "batch_size":
// validate numeric input
if value != "" {
if intVal, err := strconv.Atoi(value); err != nil || intVal <= 0 || intVal > 10000 {
return fmt.Errorf("invalid batch size: %s (must be a number between 1 and 10000)", value)
}
}
newConfig.BatchSize.Value = value
case "strip_newlines":
// validate boolean input
if value != "" && value != "true" && value != "false" {
return fmt.Errorf("invalid boolean value for strip newlines: %s (must be 'true' or 'false')", value)
}
newConfig.StripNewLines.Value = value
case "max_text_bytes":
// validate numeric input
if value != "" {
if intVal, err := strconv.Atoi(value); err != nil || intVal <= 0 {
return fmt.Errorf("invalid max text bytes: %s (must be a positive number)", value)
}
}
newConfig.MaxTextBytes.Value = value
}
}
// save the configuration
if err := m.GetController().UpdateEmbedderConfig(newConfig); err != nil {
logger.Errorf("[EmbedderFormModel] SAVE: error updating embedder config: %v", err)
return err
}
logger.Log("[EmbedderFormModel] SAVE: success")
return nil
}
func (m *EmbedderFormModel) HandleReset() {
// reset config to defaults
config := m.GetController().ResetEmbedderConfig()
// reset provider selection
selectedProvider := m.getProviderID(config.Provider.Value)
m.GetListHelper().SelectByValue(&m.providerList, selectedProvider)
// rebuild form with reset values
m.BuildForm()
}
func (m *EmbedderFormModel) OnFieldChanged(fieldIndex int, oldValue, newValue string) {
// additional validation could be added here if needed
}
func (m *EmbedderFormModel) GetFormFields() []FormField {
return m.BaseScreen.fields
}
func (m *EmbedderFormModel) SetFormFields(fields []FormField) {
m.BaseScreen.fields = fields
}
// BaseListHandler interface implementation
func (m *EmbedderFormModel) GetList() *list.Model {
return &m.providerList
}
func (m *EmbedderFormModel) GetListDelegate() *BaseListDelegate {
return m.providerDelegate
}
func (m *EmbedderFormModel) OnListSelectionChanged(oldSelection, newSelection string) {
// rebuild form when provider changes
m.BuildForm()
}
func (m *EmbedderFormModel) GetListTitle() string {
return locale.EmbedderFormProvider
}
func (m *EmbedderFormModel) GetListDescription() string {
return locale.EmbedderFormProviderDesc
}
// Update method - handle screen-specific input
func (m *EmbedderFormModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.KeyMsg:
// handle list input first (if focused on list)
if cmd := m.HandleListInput(msg); cmd != nil {
return m, cmd
}
// then handle field input
if cmd := m.HandleFieldInput(msg); cmd != nil {
return m, cmd
}
}
// delegate to base screen for common handling
cmd := m.BaseScreen.Update(msg)
return m, cmd
}
// Compile-time interface validation
var _ BaseScreenModel = (*EmbedderFormModel)(nil)
var _ BaseScreenHandler = (*EmbedderFormModel)(nil)
var _ BaseListHandler = (*EmbedderFormModel)(nil)
backend/cmd/installer/wizard/models/eula.go
+271
View File
@@ -0,0 +1,271 @@
package models
import (
"fmt"
"pentagi/cmd/installer/files"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/viewport"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
// EULAModel represents the EULA agreement screen
type EULAModel struct {
styles styles.Styles
window window.Window
viewport viewport.Model
files files.Files
controller controller.Controller
content string
ready bool
scrolled bool
scrolledToEnd bool
}
// NewEULAModel creates a new EULA screen model
func NewEULAModel(c controller.Controller, s styles.Styles, w window.Window, f files.Files) *EULAModel {
return &EULAModel{
styles: s,
window: w,
files: f,
controller: c,
}
}
// Init implements tea.Model
func (m *EULAModel) Init() tea.Cmd {
m.resetForm()
return m.loadEULA
}
// loadEULA loads the EULA content from files
func (m *EULAModel) loadEULA() tea.Msg {
content, err := m.files.GetContent("EULA.md")
if err != nil {
logger.Errorf("[EULAModel] LOAD: file error: %v", err)
return EULALoadedMsg{
Content: fmt.Sprintf(locale.EULAErrorLoadingTitle, err),
Error: err,
}
}
rendered, err := m.styles.GetRenderer().Render(string(content))
if err != nil {
logger.Errorf("[EULAModel] LOAD: render error: %v", err)
rendered = fmt.Sprintf(locale.EULAContentFallback, string(content), err)
}
return EULALoadedMsg{
Content: rendered,
Error: nil,
}
}
func (m *EULAModel) resetForm() {
m.content = ""
m.ready = false
m.scrolled = false
m.scrolledToEnd = false
}
// Update implements tea.Model
func (m *EULAModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
m.updateViewport()
case EULALoadedMsg:
m.content = msg.Content
m.updateViewport()
m.updateScrollStatus()
return m, func() tea.Msg { return nil }
case tea.KeyMsg:
switch msg.String() {
case "y", "Y":
if m.scrolledToEnd {
logger.Log("[EULAModel] ACCEPT")
if err := m.controller.SetEulaConsent(); err != nil {
logger.Errorf("[EULAModel] CONSENT: error: %v", err)
return m, func() tea.Msg { return nil }
}
// skip eula screen write to stack and go to main menu screen straight away
return m, func() tea.Msg {
m.resetForm()
return NavigationMsg{GoBack: true, Target: MainMenuScreen}
}
}
case "n", "N":
if m.scrolledToEnd {
logger.Log("[EULAModel] REJECT")
return m, func() tea.Msg { return NavigationMsg{GoBack: true} }
}
default:
if !m.ready {
break
}
switch msg.String() {
case "enter":
m.viewport.PageDown()
case "up":
m.viewport.ScrollUp(1)
case "down":
m.viewport.ScrollDown(1)
case "left":
m.viewport.ScrollLeft(2)
case "right":
m.viewport.ScrollRight(2)
case "pgup":
m.viewport.PageUp()
case "pgdown":
m.viewport.PageDown()
case "home":
m.viewport.GotoTop()
case "end":
m.viewport.GotoBottom()
}
m.updateScrollStatus()
}
}
if m.ready {
var cmd tea.Cmd
m.viewport, cmd = m.viewport.Update(msg)
m.updateScrollStatus()
return m, cmd
}
return m, nil
}
// updateViewport sets up the viewport with proper dimensions
func (m *EULAModel) updateViewport() {
contentWidth, contentHeight := m.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 || m.content == "" {
return
}
if !m.ready {
m.viewport = viewport.New(contentWidth, contentHeight)
m.viewport.Style = lipgloss.NewStyle()
m.ready = true
} else {
m.viewport.Width = contentWidth
m.viewport.Height = contentHeight
}
m.viewport.SetContent(m.content)
m.updateScrollStatus()
}
// updateScrollStatus checks if user has scrolled to the end
func (m *EULAModel) updateScrollStatus() {
if m.ready {
m.scrolled = m.viewport.ScrollPercent() > 0
m.scrolledToEnd = m.viewport.AtBottom()
}
}
// View implements tea.Model using proper lipgloss layout
func (m *EULAModel) View() string {
contentWidth, contentHeight := m.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 {
return locale.EULALoading
}
if !m.ready || m.content == "" {
return m.renderLoading()
}
content := m.viewport.View()
return lipgloss.Place(contentWidth, contentHeight, lipgloss.Center, lipgloss.Top, content)
}
// renderLoading renders loading state
func (m *EULAModel) renderLoading() string {
contentWidth, contentHeight := m.window.GetContentSize()
loading := m.styles.Info.Render(locale.EULALoading)
return lipgloss.Place(contentWidth, contentHeight, lipgloss.Center, lipgloss.Center, loading)
}
// GetScrollInfo returns scroll information for footer display
func (m *EULAModel) GetScrollInfo() (scrolled bool, atEnd bool, percent int) {
if !m.ready {
return false, false, 0
}
percent = int(m.viewport.ScrollPercent() * 100)
return m.scrolled, m.scrolledToEnd, percent
}
// EULALoadedMsg represents successful EULA loading
type EULALoadedMsg struct {
Content string
Error error
}
// BaseScreenModel interface implementation
// GetFormTitle returns empty title for the form (glamour renders from the top)
func (m *EULAModel) GetFormTitle() string {
return ""
}
// GetFormDescription returns the description for the form (right panel)
func (m *EULAModel) GetFormDescription() string {
return locale.EULAFormDescription
}
// GetFormName returns the name for the form (right panel)
func (m *EULAModel) GetFormName() string {
return locale.EULAFormName
}
// GetFormOverview returns form overview for list screens (right panel)
func (m *EULAModel) GetFormOverview() string {
return locale.EULAFormOverview
}
// GetCurrentConfiguration returns text with current configuration for the list screens
func (m *EULAModel) GetCurrentConfiguration() string {
if m.controller.GetEulaConsent() {
return locale.EULAConfigurationAccepted
}
if m.scrolledToEnd {
return locale.EULAConfigurationRead
}
return locale.EULAConfigurationPending
}
// IsConfigured returns true if eula consent is set
func (m *EULAModel) IsConfigured() bool {
return m.controller.GetEulaConsent()
}
// GetFormHotKeys returns the hotkeys for the form (layout footer)
func (m *EULAModel) GetFormHotKeys() []string {
var hotkeys []string
if m.ready && m.content != "" {
hotkeys = append(hotkeys, "up|down", "pgup|pgdown", "home|end")
}
if m.scrolledToEnd {
hotkeys = append(hotkeys, "y|n")
}
return hotkeys
}
// Compile-time interface validation
var _ BaseScreenModel = (*EULAModel)(nil)
backend/cmd/installer/wizard/models/graphiti_form.go
+401
View File
@@ -0,0 +1,401 @@
package models
import (
"fmt"
"strings"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/list"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
const (
GraphitiURLPlaceholder = "http://graphiti:8000"
GraphitiTimeoutPlaceholder = "30"
GraphitiModelNamePlaceholder = "gpt-5-mini"
GraphitiNeo4jUserPlaceholder = "neo4j"
)
// GraphitiFormModel represents the Graphiti configuration form
type GraphitiFormModel struct {
*BaseScreen
// screen-specific components
deploymentList list.Model
deploymentDelegate *BaseListDelegate
}
// NewGraphitiFormModel creates a new Graphiti form model
func NewGraphitiFormModel(c controller.Controller, s styles.Styles, w window.Window) *GraphitiFormModel {
m := &GraphitiFormModel{}
m.BaseScreen = NewBaseScreen(c, s, w, m, m)
m.initializeDeploymentList(s)
return m
}
// initializeDeploymentList sets up the deployment type selection list
func (m *GraphitiFormModel) initializeDeploymentList(styles styles.Styles) {
options := []BaseListOption{
{Value: "embedded", Display: locale.MonitoringGraphitiEmbedded},
{Value: "external", Display: locale.MonitoringGraphitiExternal},
{Value: "disabled", Display: locale.MonitoringGraphitiDisabled},
}
m.deploymentDelegate = NewBaseListDelegate(
styles.FormLabel.Align(lipgloss.Center),
MinMenuWidth-6,
)
m.deploymentList = m.GetListHelper().CreateList(options, m.deploymentDelegate, MinMenuWidth-6, 3)
config := m.GetController().GetGraphitiConfig()
m.GetListHelper().SelectByValue(&m.deploymentList, config.DeploymentType)
}
// getSelectedDeploymentType returns the currently selected deployment type using the helper
func (m *GraphitiFormModel) getSelectedDeploymentType() string {
selectedValue := m.GetListHelper().GetSelectedValue(&m.deploymentList)
if selectedValue == "" {
return "disabled"
}
return selectedValue
}
// BaseScreenHandler interface implementation
func (m *GraphitiFormModel) BuildForm() tea.Cmd {
config := m.GetController().GetGraphitiConfig()
fields := []FormField{}
deploymentType := m.getSelectedDeploymentType()
switch deploymentType {
case "embedded":
// Embedded mode - requires timeout, model, and neo4j credentials
fields = append(fields, m.createTextField(config, "timeout",
locale.MonitoringGraphitiTimeout, locale.MonitoringGraphitiTimeoutDesc, false, GraphitiTimeoutPlaceholder,
))
fields = append(fields, m.createTextField(config, "model_name",
locale.MonitoringGraphitiModelName, locale.MonitoringGraphitiModelNameDesc, false, GraphitiModelNamePlaceholder,
))
fields = append(fields, m.createTextField(config, "neo4j_user",
locale.MonitoringGraphitiNeo4jUser, locale.MonitoringGraphitiNeo4jUserDesc, false, GraphitiNeo4jUserPlaceholder,
))
fields = append(fields, m.createTextField(config, "neo4j_password",
locale.MonitoringGraphitiNeo4jPassword, locale.MonitoringGraphitiNeo4jPasswordDesc, true, "",
))
fields = append(fields, m.createTextField(config, "neo4j_database",
locale.MonitoringGraphitiNeo4jDatabase, locale.MonitoringGraphitiNeo4jDatabaseDesc, false, GraphitiNeo4jUserPlaceholder,
))
case "external":
// External mode - requires connection details only
fields = append(fields, m.createTextField(config, "url",
locale.MonitoringGraphitiURL, locale.MonitoringGraphitiURLDesc, false, GraphitiURLPlaceholder,
))
fields = append(fields, m.createTextField(config, "timeout",
locale.MonitoringGraphitiTimeout, locale.MonitoringGraphitiTimeoutDesc, false, GraphitiTimeoutPlaceholder,
))
case "disabled":
// Disabled mode has no additional fields
}
m.SetFormFields(fields)
return nil
}
func (m *GraphitiFormModel) createTextField(
config *controller.GraphitiConfig, key, title, description string, masked bool, placeholder string,
) FormField {
var envVar loader.EnvVar
switch key {
case "url":
envVar = config.GraphitiURL
case "timeout":
envVar = config.Timeout
case "model_name":
envVar = config.ModelName
case "neo4j_user":
envVar = config.Neo4jUser
case "neo4j_password":
envVar = config.Neo4jPassword
case "neo4j_database":
envVar = config.Neo4jDatabase
}
input := NewTextInput(m.GetStyles(), m.GetWindow(), envVar)
if placeholder != "" {
input.Placeholder = placeholder
}
return FormField{
Key: key,
Title: title,
Description: description,
Required: false,
Masked: masked,
Input: input,
Value: input.Value(),
}
}
func (m *GraphitiFormModel) GetFormTitle() string {
return locale.MonitoringGraphitiFormTitle
}
func (m *GraphitiFormModel) GetFormDescription() string {
return locale.MonitoringGraphitiFormDescription
}
func (m *GraphitiFormModel) GetFormName() string {
return locale.MonitoringGraphitiFormName
}
func (m *GraphitiFormModel) GetFormSummary() string {
return ""
}
func (m *GraphitiFormModel) GetFormOverview() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.MonitoringGraphitiFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.MonitoringGraphitiFormDescription))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Render(locale.MonitoringGraphitiFormOverview))
return strings.Join(sections, "\n")
}
func (m *GraphitiFormModel) GetCurrentConfiguration() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(m.GetFormName()))
config := m.GetController().GetGraphitiConfig()
getMaskedValue := func(value string) string {
maskedValue := strings.Repeat("*", len(value))
if len(value) > 15 {
maskedValue = maskedValue[:15] + "..."
}
return maskedValue
}
switch config.DeploymentType {
case "embedded":
sections = append(sections, "• "+locale.UIMode+m.GetStyles().Success.Render(locale.MonitoringGraphitiEmbedded))
if config.GraphitiURL.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiURL, m.GetStyles().Info.Render(config.GraphitiURL.Value)))
}
if timeout := config.Timeout.Value; timeout != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiTimeout, m.GetStyles().Info.Render(timeout)))
} else if timeout := config.Timeout.Default; timeout != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiTimeout, m.GetStyles().Muted.Render(timeout)))
}
if modelName := config.ModelName.Value; modelName != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiModelName, m.GetStyles().Info.Render(modelName)))
} else if modelName := config.ModelName.Default; modelName != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiModelName, m.GetStyles().Muted.Render(modelName)))
}
if neo4jUser := config.Neo4jUser.Value; neo4jUser != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiNeo4jUser, m.GetStyles().Info.Render(neo4jUser)))
} else if neo4jUser := config.Neo4jUser.Default; neo4jUser != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiNeo4jUser, m.GetStyles().Muted.Render(neo4jUser)))
}
if neo4jPassword := config.Neo4jPassword.Value; neo4jPassword != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiNeo4jPassword, m.GetStyles().Muted.Render(getMaskedValue(neo4jPassword))))
}
if neo4jDatabase := config.Neo4jDatabase.Value; neo4jDatabase != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiNeo4jDatabase, m.GetStyles().Info.Render(neo4jDatabase)))
} else if neo4jDatabase := config.Neo4jDatabase.Default; neo4jDatabase != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiNeo4jDatabase, m.GetStyles().Muted.Render(neo4jDatabase)))
}
case "external":
sections = append(sections, "• "+locale.UIMode+m.GetStyles().Success.Render(locale.MonitoringGraphitiExternal))
if config.GraphitiURL.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiURL, m.GetStyles().Info.Render(config.GraphitiURL.Value)))
}
if timeout := config.Timeout.Value; timeout != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiTimeout, m.GetStyles().Info.Render(timeout)))
} else if timeout := config.Timeout.Default; timeout != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringGraphitiTimeout, m.GetStyles().Muted.Render(timeout)))
}
case "disabled":
sections = append(sections, "• "+locale.UIMode+m.GetStyles().Warning.Render(locale.MonitoringGraphitiDisabled))
}
return strings.Join(sections, "\n")
}
func (m *GraphitiFormModel) IsConfigured() bool {
config := m.GetController().GetGraphitiConfig()
return config.DeploymentType != "disabled"
}
func (m *GraphitiFormModel) GetHelpContent() string {
var sections []string
deploymentType := m.getSelectedDeploymentType()
sections = append(sections, m.GetStyles().Subtitle.Render(locale.MonitoringGraphitiFormTitle))
sections = append(sections, "")
sections = append(sections, locale.MonitoringGraphitiModeGuide)
sections = append(sections, "")
switch deploymentType {
case "embedded":
sections = append(sections, locale.MonitoringGraphitiEmbeddedHelp)
case "external":
sections = append(sections, locale.MonitoringGraphitiExternalHelp)
case "disabled":
sections = append(sections, locale.MonitoringGraphitiDisabledHelp)
}
return strings.Join(sections, "\n")
}
func (m *GraphitiFormModel) HandleSave() error {
config := m.GetController().GetGraphitiConfig()
deploymentType := m.getSelectedDeploymentType()
fields := m.GetFormFields()
// create a working copy of the current config to modify
newConfig := &controller.GraphitiConfig{
DeploymentType: deploymentType,
// copy current EnvVar fields - they preserve metadata like Line, IsPresent, etc.
GraphitiURL: config.GraphitiURL,
Timeout: config.Timeout,
ModelName: config.ModelName,
Neo4jUser: config.Neo4jUser,
Neo4jPassword: config.Neo4jPassword,
Neo4jDatabase: config.Neo4jDatabase,
Neo4jURI: config.Neo4jURI,
Installed: config.Installed,
}
// update field values based on form input
for _, field := range fields {
value := strings.TrimSpace(field.Input.Value())
switch field.Key {
case "url":
newConfig.GraphitiURL.Value = value
case "timeout":
newConfig.Timeout.Value = value
case "model_name":
newConfig.ModelName.Value = value
case "neo4j_user":
newConfig.Neo4jUser.Value = value
case "neo4j_password":
newConfig.Neo4jPassword.Value = value
case "neo4j_database":
newConfig.Neo4jDatabase.Value = value
}
}
// save the configuration
if err := m.GetController().UpdateGraphitiConfig(newConfig); err != nil {
logger.Errorf("[GraphitiFormModel] SAVE: error updating graphiti config: %v", err)
return err
}
logger.Log("[GraphitiFormModel] SAVE: success")
return nil
}
func (m *GraphitiFormModel) HandleReset() {
// reset config to defaults
config := m.GetController().ResetGraphitiConfig()
// reset deployment selection
m.GetListHelper().SelectByValue(&m.deploymentList, config.DeploymentType)
// rebuild form with reset deployment type
m.BuildForm()
}
func (m *GraphitiFormModel) OnFieldChanged(fieldIndex int, oldValue, newValue string) {
// additional validation could be added here if needed
}
func (m *GraphitiFormModel) GetFormFields() []FormField {
return m.BaseScreen.fields
}
func (m *GraphitiFormModel) SetFormFields(fields []FormField) {
m.BaseScreen.fields = fields
}
// BaseListHandler interface implementation
func (m *GraphitiFormModel) GetList() *list.Model {
return &m.deploymentList
}
func (m *GraphitiFormModel) GetListDelegate() *BaseListDelegate {
return m.deploymentDelegate
}
func (m *GraphitiFormModel) OnListSelectionChanged(oldSelection, newSelection string) {
// rebuild form when deployment type changes
m.BuildForm()
}
func (m *GraphitiFormModel) GetListTitle() string {
return locale.MonitoringGraphitiDeploymentType
}
func (m *GraphitiFormModel) GetListDescription() string {
return locale.MonitoringGraphitiDeploymentTypeDesc
}
// Update method - handle screen-specific input
func (m *GraphitiFormModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.KeyMsg:
// handle list input first (if focused on list)
if cmd := m.HandleListInput(msg); cmd != nil {
return m, cmd
}
// then handle field input
if cmd := m.HandleFieldInput(msg); cmd != nil {
return m, cmd
}
}
// delegate to base screen for common handling
cmd := m.BaseScreen.Update(msg)
return m, cmd
}
// Compile-time interface validation
var _ BaseScreenModel = (*GraphitiFormModel)(nil)
var _ BaseScreenHandler = (*GraphitiFormModel)(nil)
var _ BaseListHandler = (*GraphitiFormModel)(nil)
backend/cmd/installer/wizard/models/helpers/calc_context.go
+307
View File
@@ -0,0 +1,307 @@
package helpers
import (
"pentagi/pkg/csum"
)
// ContextEstimate represents the estimated context size range
type ContextEstimate struct {
MinTokens int // Minimum estimated tokens (optimal summarization)
MaxTokens int // Maximum estimated tokens (approaching limits)
MinBytes int // Minimum estimated bytes
MaxBytes int // Maximum estimated bytes
}
// Global parameter boundaries based on data nature and algorithm constraints
var (
// Absolute minimums based on data nature
MinBodyPairBytes = 512 // Minimum possible body pair size
MinSectionBytes = 3 * 1024 // Minimum section: header + 1 body pair
MinSystemMessageBytes = 1 * 1024 // Minimum system message
MinHumanMessageBytes = 512 // Minimum human message
MinQASections = 1 // At least one QA section
MinKeepQASections = 1 // At least one section to keep
// Typical sizes for realistic estimation
TypicalSystemMessageBytes = 4 * 1024 // ~1k tokens for system message
TypicalHumanMessageBytes = 2 * 1024 // ~512 tokens for human message
TypicalBodyPairBytes = 8 * 1024 // ~2k tokens typical body pair
SummarizedBodyPairBytes = 6 * 1024 // ~1.5k tokens after summarization
QASummaryHeaderBytes = 4 * 1024 // ~1k tokens for QA summary header
QASummaryBodyPairBytes = 8 * 1024 // ~2k tokens for QA summary body pair
// Reasonable ranges where parameters still have meaningful impact
ReasonableMinBodyPairBytes = 8 * 1024 // 8KB
ReasonableMaxBodyPairBytes = 32 * 1024 // 32KB
ReasonableMinSectionBytes = 15 * 1024 // 15KB
ReasonableMaxSectionBytes = 100 * 1024 // 100KB
ReasonableMinQABytes = 30 * 1024 // 30KB
ReasonableMaxQABytes = 500 * 1024 // 500KB
ReasonableMinQASections = 2 // 2 sections
ReasonableMaxQASections = 15 // 15 sections
// Token to byte conversion ratio
TokenToByteRatio = 4
)
// ConfigBoundaries represents effective boundaries for a specific configuration
type ConfigBoundaries struct {
// Effective ranges for parameters based on configuration
MinBodyPairBytes int
MaxBodyPairBytes int
MinSectionBytes int
MaxSectionBytes int
MinQABytes int
MaxQABytes int
MinQASections int
MaxQASections int
MinKeepSections int
MaxKeepSections int
// Derived boundaries
MinSectionsToProcess int // Minimum sections that would be processed
MaxSectionsToProcess int // Maximum sections before QA triggers
}
// NewConfigBoundaries creates boundaries adjusted for specific configuration
func NewConfigBoundaries(config csum.SummarizerConfig) ConfigBoundaries {
boundaries := ConfigBoundaries{}
// Body pair boundaries
boundaries.MinBodyPairBytes = max(MinBodyPairBytes, ReasonableMinBodyPairBytes)
if config.MaxBPBytes > 0 {
boundaries.MaxBodyPairBytes = min(config.MaxBPBytes, ReasonableMaxBodyPairBytes)
} else {
boundaries.MaxBodyPairBytes = ReasonableMaxBodyPairBytes
}
boundaries.MaxBodyPairBytes = max(boundaries.MaxBodyPairBytes, boundaries.MinBodyPairBytes)
// Section boundaries
boundaries.MinSectionBytes = max(MinSectionBytes, ReasonableMinSectionBytes)
if config.LastSecBytes > 0 {
boundaries.MaxSectionBytes = min(config.LastSecBytes, ReasonableMaxSectionBytes)
} else {
boundaries.MaxSectionBytes = ReasonableMaxSectionBytes
}
boundaries.MaxSectionBytes = max(boundaries.MaxSectionBytes, boundaries.MinSectionBytes)
// QA bytes boundaries
boundaries.MinQABytes = max(boundaries.MinSectionBytes*ReasonableMinQASections, ReasonableMinQABytes)
if config.MaxQABytes > 0 {
boundaries.MaxQABytes = min(config.MaxQABytes, ReasonableMaxQABytes)
} else {
boundaries.MaxQABytes = ReasonableMaxQABytes
}
boundaries.MaxQABytes = max(boundaries.MaxQABytes, boundaries.MinQABytes)
// QA sections boundaries
boundaries.MinQASections = max(MinQASections, ReasonableMinQASections)
if config.MaxQASections > 0 {
boundaries.MaxQASections = min(config.MaxQASections, ReasonableMaxQASections)
} else {
boundaries.MaxQASections = ReasonableMaxQASections
}
boundaries.MaxQASections = max(boundaries.MaxQASections, boundaries.MinQASections)
// Keep sections boundaries
boundaries.MinKeepSections = max(MinKeepQASections, config.KeepQASections)
boundaries.MaxKeepSections = min(boundaries.MaxQASections, config.KeepQASections)
boundaries.MaxKeepSections = max(boundaries.MaxKeepSections, boundaries.MinKeepSections)
// Derived boundaries for sections processing
boundaries.MinSectionsToProcess = boundaries.MinKeepSections + 1
// Calculate when QA summarization would trigger
minSectionSize := boundaries.MinSectionBytes
maxSectionsBeforeQA := boundaries.MaxQABytes / minSectionSize
boundaries.MaxSectionsToProcess = min(maxSectionsBeforeQA, boundaries.MaxQASections)
boundaries.MaxSectionsToProcess = max(boundaries.MaxSectionsToProcess, boundaries.MinSectionsToProcess)
return boundaries
}
// CalculateContextEstimate calculates the estimated context size based on summarizer configuration
func CalculateContextEstimate(config csum.SummarizerConfig) ContextEstimate {
// Create boundaries for this configuration
boundaries := NewConfigBoundaries(config)
// Calculate minimum context (optimal summarization scenario)
minBytes := calculateMinimumContext(config, boundaries)
// Calculate maximum context (approaching limits scenario)
maxBytes := calculateMaximumContext(config, boundaries)
// Convert bytes to tokens
minTokens := minBytes / TokenToByteRatio
maxTokens := maxBytes / TokenToByteRatio
return ContextEstimate{
MinTokens: minTokens,
MaxTokens: maxTokens,
MinBytes: minBytes,
MaxBytes: maxBytes,
}
}
// calculateMinimumContext estimates minimum context when summarization works optimally
func calculateMinimumContext(config csum.SummarizerConfig, boundaries ConfigBoundaries) int {
totalBytes := 0
// Base overhead: system message
totalBytes += TypicalSystemMessageBytes
// Base sections: use boundaries for minimum sections count
baseSections := max(boundaries.MinKeepSections, 1)
// Use boundaries for minimum section size calculation
minSectionSize := TypicalHumanMessageBytes + SummarizedBodyPairBytes
// For each base section, calculate minimal content
for i := 0; i < baseSections; i++ {
// Section header
totalBytes += TypicalHumanMessageBytes
// Section body - use minimal body pair sizes from boundaries
var sectionBodySize int
if config.PreserveLast {
// PreserveLast=true means sections are summarized to fit LastSecBytes
// This REDUCES total size (more compression)
sectionBodySize = min(boundaries.MinSectionBytes/3, SummarizedBodyPairBytes*2)
} else {
// PreserveLast=false means sections remain as-is without last section management
// This INCREASES total size (less compression)
sectionBodySize = min(boundaries.MinSectionBytes, TypicalBodyPairBytes*2)
}
totalBytes += sectionBodySize
}
// Add minimal impact from configuration parameters using boundaries
// MaxBPBytes influence through boundaries
if boundaries.MaxBodyPairBytes > boundaries.MinBodyPairBytes {
// Add small fraction of the difference for minimal scenario
overhead := (boundaries.MaxBodyPairBytes - boundaries.MinBodyPairBytes) / 10
totalBytes += overhead
}
// MaxQABytes influence through boundaries
if boundaries.MaxQABytes > boundaries.MinQABytes {
// Add small fraction for potential QA content
qaOverhead := (boundaries.MaxQABytes - boundaries.MinQABytes) / 20
totalBytes += qaOverhead
}
// MaxQASections influence through boundaries
if boundaries.MaxQASections > boundaries.MinQASections {
// Each additional section beyond minimum adds fractional content
additionalSections := boundaries.MaxQASections - boundaries.MinQASections
totalBytes += min(additionalSections, 3) * (minSectionSize / 4)
}
// Boolean parameter influences for optimization
if config.UseQA {
// QA summarization generally reduces context through better summarization
totalBytes = totalBytes * 9 / 10 // 10% reduction
if config.SummHumanInQA {
// Summarizing human messages saves additional space
totalBytes = totalBytes * 95 / 100 // Additional 5% reduction
}
}
// Note: PreserveLast effect is already built into section calculation above
return totalBytes
}
// calculateMaximumContext estimates maximum context when approaching configuration limits
func calculateMaximumContext(config csum.SummarizerConfig, boundaries ConfigBoundaries) int {
totalBytes := 0
// Base overhead: system message
totalBytes += TypicalSystemMessageBytes
// Base sections: use boundaries for maximum sections that would be processed
baseSections := max(boundaries.MaxKeepSections, boundaries.MinKeepSections)
// Use boundaries for maximum section size calculation
maxSectionSize := TypicalHumanMessageBytes + TypicalBodyPairBytes
// For each base section, calculate maximum content
for i := 0; i < baseSections; i++ {
// Section header
totalBytes += TypicalHumanMessageBytes
// Section body - use maximum body pair sizes from boundaries
var sectionBodySize int
if config.PreserveLast {
// PreserveLast=true means sections are summarized to fit within LastSecBytes
// This keeps size SMALLER (more compression)
sectionBodySize = min(boundaries.MaxSectionBytes, TypicalBodyPairBytes*3)
} else {
// PreserveLast=false means sections can grow larger without management
// This allows size to be LARGER (less compression)
sectionBodySize = min(boundaries.MaxSectionBytes*2, ReasonableMaxSectionBytes)
}
totalBytes += sectionBodySize
}
// Add maximum impact from configuration parameters using boundaries
// MaxBPBytes influence through boundaries - full impact in maximum scenario
if boundaries.MaxBodyPairBytes > boundaries.MinBodyPairBytes {
// Add significant portion of the difference for maximum scenario
overhead := (boundaries.MaxBodyPairBytes - boundaries.MinBodyPairBytes) / 2
totalBytes += overhead
}
// MaxQABytes influence through boundaries - substantial impact
if boundaries.MaxQABytes > boundaries.MinQABytes {
// Add larger fraction for maximum QA content potential
qaOverhead := (boundaries.MaxQABytes - boundaries.MinQABytes) / 8
totalBytes += qaOverhead
}
// MaxQASections influence through boundaries - linear growth
if boundaries.MaxQASections > boundaries.MinQASections {
// Each additional section beyond minimum adds substantial content in max scenario
additionalSections := boundaries.MaxQASections - boundaries.MinQASections
totalBytes += min(additionalSections, 8) * (maxSectionSize / 2) // Half section size per additional
}
// Boolean parameter influences - mainly affecting complexity/growth
if config.UseQA {
// QA can enable more complex conversations in maximum scenario
totalBytes = totalBytes * 105 / 100 // 5% increase for QA complexity
if config.SummHumanInQA {
// Summarizing human messages reduces maximum size
totalBytes = totalBytes * 95 / 100 // 5% reduction
}
} else {
// Without QA, conversations can be less organized and grow larger
totalBytes = totalBytes * 110 / 100 // 10% increase
}
// KeepQASections direct influence: more kept sections = linearly more content
// Use boundaries to ensure consistency
keepSections := max(boundaries.MaxKeepSections, 1)
if keepSections > 1 {
// Each additional kept section adds substantial content in max scenario
additionalKeptSections := keepSections - 1
additionalContent := additionalKeptSections * maxSectionSize
// Apply PreserveLast effect to additional content
if config.PreserveLast {
// PreserveLast reduces the size of additional content
additionalContent = additionalContent * 8 / 10 // 20% reduction
}
totalBytes += additionalContent
}
// Add small buffer for message structure overhead
totalBytes += 2 * 1024
return totalBytes
}
backend/cmd/installer/wizard/models/helpers/calc_context_test.go
+507
View File
@@ -0,0 +1,507 @@
package helpers
import (
"testing"
"pentagi/pkg/csum"
)
// TestConfigBoundaries verifies that boundaries are calculated correctly based on configuration
func TestConfigBoundaries(t *testing.T) {
tests := []struct {
name string
config csum.SummarizerConfig
verify func(t *testing.T, boundaries ConfigBoundaries)
}{
{
name: "Default configuration boundaries",
config: csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 50 * 1024,
MaxBPBytes: 16 * 1024,
MaxQABytes: 100 * 1024,
MaxQASections: 10,
KeepQASections: 2,
},
verify: func(t *testing.T, b ConfigBoundaries) {
// Check that boundaries respect configuration limits
if b.MaxSectionBytes > 50*1024 {
t.Errorf("MaxSectionBytes (%d) should not exceed LastSecBytes (50KB)", b.MaxSectionBytes)
}
if b.MaxBodyPairBytes > 16*1024 {
t.Errorf("MaxBodyPairBytes (%d) should not exceed MaxBPBytes (16KB)", b.MaxBodyPairBytes)
}
if b.MaxQABytes > 100*1024 {
t.Errorf("MaxQABytes (%d) should not exceed config MaxQABytes (100KB)", b.MaxQABytes)
}
// Check minimum bounds
if b.MinSectionBytes < MinSectionBytes {
t.Errorf("MinSectionBytes (%d) below absolute minimum (%d)", b.MinSectionBytes, MinSectionBytes)
}
if b.MinBodyPairBytes < MinBodyPairBytes {
t.Errorf("MinBodyPairBytes (%d) below absolute minimum (%d)", b.MinBodyPairBytes, MinBodyPairBytes)
}
// Check logical consistency
if b.MaxSectionBytes < b.MinSectionBytes {
t.Errorf("MaxSectionBytes (%d) < MinSectionBytes (%d)", b.MaxSectionBytes, b.MinSectionBytes)
}
if b.MaxBodyPairBytes < b.MinBodyPairBytes {
t.Errorf("MaxBodyPairBytes (%d) < MinBodyPairBytes (%d)", b.MaxBodyPairBytes, b.MinBodyPairBytes)
}
},
},
{
name: "Extreme low values boundaries",
config: csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 10 * 1024, // Very low
MaxBPBytes: 4 * 1024, // Very low
MaxQABytes: 20 * 1024, // Very low
MaxQASections: 2, // Very low
KeepQASections: 1,
},
verify: func(t *testing.T, b ConfigBoundaries) {
// Even with low config values, boundaries should not go below reasonable minimums
if b.MinSectionBytes < ReasonableMinSectionBytes {
t.Errorf("MinSectionBytes (%d) should be at least reasonable minimum (%d)",
b.MinSectionBytes, ReasonableMinSectionBytes)
}
if b.MinBodyPairBytes < ReasonableMinBodyPairBytes {
t.Errorf("MinBodyPairBytes (%d) should be at least reasonable minimum (%d)",
b.MinBodyPairBytes, ReasonableMinBodyPairBytes)
}
},
},
{
name: "Extreme high values boundaries",
config: csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 200 * 1024, // Very high
MaxBPBytes: 64 * 1024, // Very high
MaxQABytes: 1024 * 1024, // Very high
MaxQASections: 50, // Very high
KeepQASections: 20,
},
verify: func(t *testing.T, b ConfigBoundaries) {
// High config values should be capped at reasonable maximums
if b.MaxSectionBytes > ReasonableMaxSectionBytes {
t.Errorf("MaxSectionBytes (%d) should be capped at reasonable maximum (%d)",
b.MaxSectionBytes, ReasonableMaxSectionBytes)
}
if b.MaxBodyPairBytes > ReasonableMaxBodyPairBytes {
t.Errorf("MaxBodyPairBytes (%d) should be capped at reasonable maximum (%d)",
b.MaxBodyPairBytes, ReasonableMaxBodyPairBytes)
}
if b.MaxQABytes > ReasonableMaxQABytes {
t.Errorf("MaxQABytes (%d) should be capped at reasonable maximum (%d)",
b.MaxQABytes, ReasonableMaxQABytes)
}
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
boundaries := NewConfigBoundaries(tt.config)
t.Logf("Config: LastSec=%dKB, MaxBP=%dKB, MaxQA=%dKB, MaxQASections=%d, KeepQA=%d",
tt.config.LastSecBytes/1024, tt.config.MaxBPBytes/1024, tt.config.MaxQABytes/1024,
tt.config.MaxQASections, tt.config.KeepQASections)
t.Logf("Boundaries: MinSec=%dKB, MaxSec=%dKB, MinBP=%dKB, MaxBP=%dKB, MinQA=%dKB, MaxQA=%dKB",
boundaries.MinSectionBytes/1024, boundaries.MaxSectionBytes/1024,
boundaries.MinBodyPairBytes/1024, boundaries.MaxBodyPairBytes/1024,
boundaries.MinQABytes/1024, boundaries.MaxQABytes/1024)
tt.verify(t, boundaries)
})
}
}
// TestMonotonicBehavior tests that increasing each parameter never decreases context estimates
func TestMonotonicBehavior(t *testing.T) {
baseConfig := csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 50 * 1024,
MaxBPBytes: 16 * 1024,
MaxQABytes: 100 * 1024,
MaxQASections: 10,
KeepQASections: 2,
}
// Test KeepQASections monotonicity (most important parameter)
t.Run("KeepQASections", func(t *testing.T) {
var prevEstimate *ContextEstimate
for _, keepSections := range []int{1, 2, 3, 5, 7, 10} {
config := baseConfig
config.KeepQASections = keepSections
config.MaxQASections = max(config.MaxQASections, keepSections) // Ensure consistency
estimate := CalculateContextEstimate(config)
t.Logf("KeepQASections=%d: Min=%d, Max=%d tokens",
keepSections, estimate.MinTokens, estimate.MaxTokens)
if prevEstimate != nil {
if estimate.MinTokens < prevEstimate.MinTokens {
t.Errorf("Non-monotonic MinTokens: %d < %d for KeepQASections %d",
estimate.MinTokens, prevEstimate.MinTokens, keepSections)
}
if estimate.MaxTokens < prevEstimate.MaxTokens {
t.Errorf("Non-monotonic MaxTokens: %d < %d for KeepQASections %d",
estimate.MaxTokens, prevEstimate.MaxTokens, keepSections)
}
}
prevEstimate = &estimate
}
})
// Test LastSecBytes monotonicity
t.Run("LastSecBytes", func(t *testing.T) {
var prevEstimate *ContextEstimate
for _, lastSecBytes := range []int{20 * 1024, 30 * 1024, 50 * 1024, 70 * 1024, 100 * 1024} {
config := baseConfig
config.LastSecBytes = lastSecBytes
estimate := CalculateContextEstimate(config)
t.Logf("LastSecBytes=%dKB: Min=%d, Max=%d tokens",
lastSecBytes/1024, estimate.MinTokens, estimate.MaxTokens)
if prevEstimate != nil {
if estimate.MinTokens < prevEstimate.MinTokens {
t.Errorf("Non-monotonic MinTokens: %d < %d for LastSecBytes %dKB",
estimate.MinTokens, prevEstimate.MinTokens, lastSecBytes/1024)
}
if estimate.MaxTokens < prevEstimate.MaxTokens {
t.Errorf("Non-monotonic MaxTokens: %d < %d for LastSecBytes %dKB",
estimate.MaxTokens, prevEstimate.MaxTokens, lastSecBytes/1024)
}
}
prevEstimate = &estimate
}
})
// Test MaxQABytes monotonicity
t.Run("MaxQABytes", func(t *testing.T) {
var prevEstimate *ContextEstimate
for _, maxQABytes := range []int{50 * 1024, 75 * 1024, 100 * 1024, 150 * 1024, 200 * 1024} {
config := baseConfig
config.MaxQABytes = maxQABytes
estimate := CalculateContextEstimate(config)
t.Logf("MaxQABytes=%dKB: Min=%d, Max=%d tokens",
maxQABytes/1024, estimate.MinTokens, estimate.MaxTokens)
if prevEstimate != nil {
if estimate.MinTokens < prevEstimate.MinTokens {
t.Errorf("Non-monotonic MinTokens: %d < %d for MaxQABytes %dKB",
estimate.MinTokens, prevEstimate.MinTokens, maxQABytes/1024)
}
if estimate.MaxTokens < prevEstimate.MaxTokens {
t.Errorf("Non-monotonic MaxTokens: %d < %d for MaxQABytes %dKB",
estimate.MaxTokens, prevEstimate.MaxTokens, maxQABytes/1024)
}
}
prevEstimate = &estimate
}
})
// Test MaxQASections monotonicity
t.Run("MaxQASections", func(t *testing.T) {
var prevEstimate *ContextEstimate
for _, maxQASections := range []int{3, 5, 8, 10, 15} {
config := baseConfig
config.MaxQASections = maxQASections
estimate := CalculateContextEstimate(config)
t.Logf("MaxQASections=%d: Min=%d, Max=%d tokens",
maxQASections, estimate.MinTokens, estimate.MaxTokens)
if prevEstimate != nil {
if estimate.MinTokens < prevEstimate.MinTokens {
t.Errorf("Non-monotonic MinTokens: %d < %d for MaxQASections %d",
estimate.MinTokens, prevEstimate.MinTokens, maxQASections)
}
if estimate.MaxTokens < prevEstimate.MaxTokens {
t.Errorf("Non-monotonic MaxTokens: %d < %d for MaxQASections %d",
estimate.MaxTokens, prevEstimate.MaxTokens, maxQASections)
}
}
prevEstimate = &estimate
}
})
// Test MaxBPBytes monotonicity
t.Run("MaxBPBytes", func(t *testing.T) {
var prevEstimate *ContextEstimate
for _, maxBPBytes := range []int{8 * 1024, 12 * 1024, 16 * 1024, 24 * 1024, 32 * 1024} {
config := baseConfig
config.MaxBPBytes = maxBPBytes
estimate := CalculateContextEstimate(config)
t.Logf("MaxBPBytes=%dKB: Min=%d, Max=%d tokens",
maxBPBytes/1024, estimate.MinTokens, estimate.MaxTokens)
if prevEstimate != nil {
if estimate.MinTokens < prevEstimate.MinTokens {
t.Errorf("Non-monotonic MinTokens: %d < %d for MaxBPBytes %dKB",
estimate.MinTokens, prevEstimate.MinTokens, maxBPBytes/1024)
}
if estimate.MaxTokens < prevEstimate.MaxTokens {
t.Errorf("Non-monotonic MaxTokens: %d < %d for MaxBPBytes %dKB",
estimate.MaxTokens, prevEstimate.MaxTokens, maxBPBytes/1024)
}
}
prevEstimate = &estimate
}
})
}
// TestBooleanParametersLogic tests correct behavior of boolean parameters
func TestBooleanParametersLogic(t *testing.T) {
baseConfig := csum.SummarizerConfig{
LastSecBytes: 50 * 1024,
MaxBPBytes: 16 * 1024,
MaxQABytes: 100 * 1024,
MaxQASections: 10,
KeepQASections: 3,
}
// Test PreserveLast parameter (CRITICAL TEST)
t.Run("PreserveLast", func(t *testing.T) {
configFalse := baseConfig
configFalse.PreserveLast = false
configFalse.UseQA = true
configFalse.SummHumanInQA = false
configTrue := baseConfig
configTrue.PreserveLast = true
configTrue.UseQA = true
configTrue.SummHumanInQA = false
estimateFalse := CalculateContextEstimate(configFalse)
estimateTrue := CalculateContextEstimate(configTrue)
t.Logf("PreserveLast=false: Min=%d, Max=%d tokens", estimateFalse.MinTokens, estimateFalse.MaxTokens)
t.Logf("PreserveLast=true: Min=%d, Max=%d tokens", estimateTrue.MinTokens, estimateTrue.MaxTokens)
// CRITICAL: PreserveLast=true should result in SMALLER context (more summarization)
// PreserveLast=false should result in LARGER context (less summarization)
if estimateTrue.MaxTokens >= estimateFalse.MaxTokens {
t.Errorf("PreserveLast=true should produce SMALLER MaxTokens than false. Got true=%d, false=%d",
estimateTrue.MaxTokens, estimateFalse.MaxTokens)
}
if estimateTrue.MinTokens > estimateFalse.MinTokens {
t.Errorf("PreserveLast=true should produce SMALLER or equal MinTokens than false. Got true=%d, false=%d",
estimateTrue.MinTokens, estimateFalse.MinTokens)
}
})
// Test UseQA parameter
t.Run("UseQA", func(t *testing.T) {
configFalse := baseConfig
configFalse.PreserveLast = true
configFalse.UseQA = false
configFalse.SummHumanInQA = false
configTrue := baseConfig
configTrue.PreserveLast = true
configTrue.UseQA = true
configTrue.SummHumanInQA = false
estimateFalse := CalculateContextEstimate(configFalse)
estimateTrue := CalculateContextEstimate(configTrue)
t.Logf("UseQA=false: Min=%d, Max=%d tokens", estimateFalse.MinTokens, estimateFalse.MaxTokens)
t.Logf("UseQA=true: Min=%d, Max=%d tokens", estimateTrue.MinTokens, estimateTrue.MaxTokens)
// UseQA should affect the results (direction depends on scenario, but should be different)
if estimateFalse.MinTokens == estimateTrue.MinTokens && estimateFalse.MaxTokens == estimateTrue.MaxTokens {
t.Errorf("UseQA parameter should affect the estimates")
}
})
// Test SummHumanInQA parameter
t.Run("SummHumanInQA", func(t *testing.T) {
configFalse := baseConfig
configFalse.PreserveLast = true
configFalse.UseQA = true
configFalse.SummHumanInQA = false
configTrue := baseConfig
configTrue.PreserveLast = true
configTrue.UseQA = true
configTrue.SummHumanInQA = true
estimateFalse := CalculateContextEstimate(configFalse)
estimateTrue := CalculateContextEstimate(configTrue)
t.Logf("SummHumanInQA=false: Min=%d, Max=%d tokens", estimateFalse.MinTokens, estimateFalse.MaxTokens)
t.Logf("SummHumanInQA=true: Min=%d, Max=%d tokens", estimateTrue.MinTokens, estimateTrue.MaxTokens)
// SummHumanInQA=true should result in smaller context (more summarization)
if estimateTrue.MaxTokens > estimateFalse.MaxTokens {
t.Errorf("SummHumanInQA=true should produce smaller or equal MaxTokens than false. Got true=%d, false=%d",
estimateTrue.MaxTokens, estimateFalse.MaxTokens)
}
})
}
// TestBoundariesUsage verifies that calculation functions actually use the boundaries
func TestBoundariesUsage(t *testing.T) {
// This test ensures that boundaries are actually used in calculations
config1 := csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 30 * 1024, // Low value
MaxBPBytes: 8 * 1024, // Low value
MaxQABytes: 50 * 1024, // Low value
MaxQASections: 5,
KeepQASections: 2,
}
config2 := csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 80 * 1024, // High value
MaxBPBytes: 24 * 1024, // High value
MaxQABytes: 200 * 1024, // High value
MaxQASections: 5, // Same as config1
KeepQASections: 2, // Same as config1
}
boundaries1 := NewConfigBoundaries(config1)
boundaries2 := NewConfigBoundaries(config2)
// Boundaries should be different
if boundaries1.MaxSectionBytes == boundaries2.MaxSectionBytes {
t.Errorf("Boundaries should differ based on configuration")
}
estimate1 := CalculateContextEstimate(config1)
estimate2 := CalculateContextEstimate(config2)
t.Logf("Config1 boundaries: MaxSec=%dKB, MaxBP=%dKB, MaxQA=%dKB",
boundaries1.MaxSectionBytes/1024, boundaries1.MaxBodyPairBytes/1024, boundaries1.MaxQABytes/1024)
t.Logf("Config2 boundaries: MaxSec=%dKB, MaxBP=%dKB, MaxQA=%dKB",
boundaries2.MaxSectionBytes/1024, boundaries2.MaxBodyPairBytes/1024, boundaries2.MaxQABytes/1024)
t.Logf("Config1 estimate: Min=%d, Max=%d tokens", estimate1.MinTokens, estimate1.MaxTokens)
t.Logf("Config2 estimate: Min=%d, Max=%d tokens", estimate2.MinTokens, estimate2.MaxTokens)
// Config2 should have larger estimates since it has larger limits
if estimate2.MaxTokens <= estimate1.MaxTokens {
t.Errorf("Config with larger limits should produce larger estimates. Got config1=%d, config2=%d",
estimate1.MaxTokens, estimate2.MaxTokens)
}
}
// TestCalculateContextEstimate verifies the main function works correctly
func TestCalculateContextEstimate(t *testing.T) {
testCases := []struct {
name string
config csum.SummarizerConfig
verify func(t *testing.T, estimate ContextEstimate)
}{
{
name: "Minimal configuration",
config: csum.SummarizerConfig{
PreserveLast: false,
UseQA: false,
SummHumanInQA: false,
LastSecBytes: 20 * 1024,
MaxBPBytes: 8 * 1024,
MaxQABytes: 30 * 1024,
MaxQASections: 3,
KeepQASections: 1,
},
verify: func(t *testing.T, estimate ContextEstimate) {
if estimate.MinTokens <= 0 || estimate.MaxTokens <= estimate.MinTokens {
t.Errorf("Invalid estimates: Min=%d, Max=%d", estimate.MinTokens, estimate.MaxTokens)
}
// Should be relatively small
if estimate.MaxTokens > 20000 {
t.Errorf("Minimal config should produce modest estimates, got %d tokens", estimate.MaxTokens)
}
},
},
{
name: "Maximal configuration",
config: csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: true,
LastSecBytes: 100 * 1024,
MaxBPBytes: 32 * 1024,
MaxQABytes: 400 * 1024,
MaxQASections: 15,
KeepQASections: 8,
},
verify: func(t *testing.T, estimate ContextEstimate) {
if estimate.MinTokens <= 0 || estimate.MaxTokens <= estimate.MinTokens {
t.Errorf("Invalid estimates: Min=%d, Max=%d", estimate.MinTokens, estimate.MaxTokens)
}
// Should be larger than minimal config
if estimate.MaxTokens < 30000 {
t.Errorf("Maximal config should produce substantial estimates, got %d tokens", estimate.MaxTokens)
}
},
},
{
name: "Default-like configuration",
config: csum.SummarizerConfig{
PreserveLast: true,
UseQA: true,
SummHumanInQA: false,
LastSecBytes: 50 * 1024,
MaxBPBytes: 16 * 1024,
MaxQABytes: 64 * 1024,
MaxQASections: 10,
KeepQASections: 1,
},
verify: func(t *testing.T, estimate ContextEstimate) {
// Check token to byte ratio
expectedMinBytes := estimate.MinTokens * TokenToByteRatio
expectedMaxBytes := estimate.MaxTokens * TokenToByteRatio
// Allow small rounding errors (due to divisions in calculation)
if abs(estimate.MinBytes-expectedMinBytes) > 5 {
t.Errorf("MinBytes calculation error: expected %d, got %d", expectedMinBytes, estimate.MinBytes)
}
if abs(estimate.MaxBytes-expectedMaxBytes) > 5 {
t.Errorf("MaxBytes calculation error: expected %d, got %d", expectedMaxBytes, estimate.MaxBytes)
}
},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
estimate := CalculateContextEstimate(tc.config)
t.Logf("%s: Min=%d tokens (%d bytes), Max=%d tokens (%d bytes)",
tc.name, estimate.MinTokens, estimate.MinBytes, estimate.MaxTokens, estimate.MaxBytes)
tc.verify(t, estimate)
})
}
}
// Helper function for absolute value
func abs(x int) int {
if x < 0 {
return -x
}
return x
}
backend/cmd/installer/wizard/models/langfuse_form.go
+456
View File
@@ -0,0 +1,456 @@
package models
import (
"fmt"
"strings"
"pentagi/cmd/installer/loader"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/logger"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/list"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
const (
LangfuseBaseURLPlaceholder = "https://cloud.langfuse.com"
LangfuseProjectIDPlaceholder = "cm000000000000000000000000"
LangfusePublicKeyPlaceholder = "pk-lf-00000000-0000-0000-0000-000000000000"
LangfuseSecretKeyPlaceholder = ""
LangfuseAdminEmailPlaceholder = "admin@pentagi.com"
LangfuseAdminPasswordPlaceholder = ""
LangfuseAdminNamePlaceholder = "admin"
LangfuseLicenseKeyPlaceholder = "sk-lf-ee-xxxxxxxxxxxxxxxxxxxxxxxx"
)
// LangfuseFormModel represents the Langfuse configuration form
type LangfuseFormModel struct {
*BaseScreen
// screen-specific components
deploymentList list.Model
deploymentDelegate *BaseListDelegate
}
// NewLangfuseFormModel creates a new Langfuse form model
func NewLangfuseFormModel(c controller.Controller, s styles.Styles, w window.Window) *LangfuseFormModel {
m := &LangfuseFormModel{}
m.BaseScreen = NewBaseScreen(c, s, w, m, m)
m.initializeDeploymentList(s)
return m
}
// initializeDeploymentList sets up the deployment type selection list
func (m *LangfuseFormModel) initializeDeploymentList(styles styles.Styles) {
options := []BaseListOption{
{Value: "embedded", Display: locale.MonitoringLangfuseEmbedded},
{Value: "external", Display: locale.MonitoringLangfuseExternal},
{Value: "disabled", Display: locale.MonitoringLangfuseDisabled},
}
m.deploymentDelegate = NewBaseListDelegate(
styles.FormLabel.Align(lipgloss.Center),
MinMenuWidth-6,
)
m.deploymentList = m.GetListHelper().CreateList(options, m.deploymentDelegate, MinMenuWidth-6, 3)
config := m.GetController().GetLangfuseConfig()
m.GetListHelper().SelectByValue(&m.deploymentList, config.DeploymentType)
}
// getSelectedDeploymentType returns the currently selected deployment type using the helper
func (m *LangfuseFormModel) getSelectedDeploymentType() string {
selectedValue := m.GetListHelper().GetSelectedValue(&m.deploymentList)
if selectedValue == "" {
return "disabled"
}
return selectedValue
}
// BaseScreenHandler interface implementation
func (m *LangfuseFormModel) BuildForm() tea.Cmd {
config := m.GetController().GetLangfuseConfig()
fields := []FormField{}
deploymentType := m.getSelectedDeploymentType()
switch deploymentType {
case "embedded":
// Embedded mode - requires all fields including admin credentials
fields = append(fields, m.createTextField(config, "listen_ip",
locale.MonitoringLangfuseListenIP, locale.MonitoringLangfuseListenIPDesc, false, "",
))
fields = append(fields, m.createTextField(config, "listen_port",
locale.MonitoringLangfuseListenPort, locale.MonitoringLangfuseListenPortDesc, false, "",
))
fields = append(fields, m.createTextField(config, "project_id",
locale.MonitoringLangfuseProjectID, locale.MonitoringLangfuseProjectIDDesc, false, LangfuseProjectIDPlaceholder,
))
fields = append(fields, m.createTextField(config, "public_key",
locale.MonitoringLangfusePublicKey, locale.MonitoringLangfusePublicKeyDesc, true, LangfusePublicKeyPlaceholder,
))
fields = append(fields, m.createTextField(config, "secret_key",
locale.MonitoringLangfuseSecretKey, locale.MonitoringLangfuseSecretKeyDesc, true, LangfuseSecretKeyPlaceholder,
))
if !config.Installed {
fields = append(fields, m.createTextField(config, "admin_email",
locale.MonitoringLangfuseAdminEmail, locale.MonitoringLangfuseAdminEmailDesc, false, LangfuseAdminEmailPlaceholder,
))
fields = append(fields, m.createTextField(config, "admin_password",
locale.MonitoringLangfuseAdminPassword, locale.MonitoringLangfuseAdminPasswordDesc, true, LangfuseAdminPasswordPlaceholder,
))
fields = append(fields, m.createTextField(config, "admin_name",
locale.MonitoringLangfuseAdminName, locale.MonitoringLangfuseAdminNameDesc, false, LangfuseAdminNamePlaceholder,
))
}
fields = append(fields, m.createTextField(config, "license_key",
locale.MonitoringLangfuseLicenseKey, locale.MonitoringLangfuseLicenseKeyDesc, true, LangfuseLicenseKeyPlaceholder,
))
case "external":
// External mode - requires connection details only
fields = append(fields, m.createTextField(config, "base_url",
locale.MonitoringLangfuseBaseURL, locale.MonitoringLangfuseBaseURLDesc, false, LangfuseBaseURLPlaceholder,
))
fields = append(fields, m.createTextField(config, "project_id",
locale.MonitoringLangfuseProjectID, locale.MonitoringLangfuseProjectIDDesc, false, LangfuseProjectIDPlaceholder,
))
fields = append(fields, m.createTextField(config, "public_key",
locale.MonitoringLangfusePublicKey, locale.MonitoringLangfusePublicKeyDesc, true, LangfusePublicKeyPlaceholder,
))
fields = append(fields, m.createTextField(config, "secret_key",
locale.MonitoringLangfuseSecretKey, locale.MonitoringLangfuseSecretKeyDesc, true, LangfuseSecretKeyPlaceholder,
))
case "disabled":
// Disabled mode has no additional fields
}
m.SetFormFields(fields)
return nil
}
func (m *LangfuseFormModel) createTextField(
config *controller.LangfuseConfig, key, title, description string, masked bool, placeholder string,
) FormField {
var envVar loader.EnvVar
switch key {
case "listen_ip":
envVar = config.ListenIP
case "listen_port":
envVar = config.ListenPort
case "base_url":
envVar = config.BaseURL
case "project_id":
envVar = config.ProjectID
case "public_key":
envVar = config.PublicKey
case "secret_key":
envVar = config.SecretKey
case "admin_email":
envVar = config.AdminEmail
case "admin_password":
envVar = config.AdminPassword
case "admin_name":
envVar = config.AdminName
case "license_key":
envVar = config.LicenseKey
}
input := NewTextInput(m.GetStyles(), m.GetWindow(), envVar)
if placeholder != "" {
input.Placeholder = placeholder
}
return FormField{
Key: key,
Title: title,
Description: description,
Required: false,
Masked: masked,
Input: input,
Value: input.Value(),
}
}
func (m *LangfuseFormModel) GetFormTitle() string {
return locale.MonitoringLangfuseFormTitle
}
func (m *LangfuseFormModel) GetFormDescription() string {
return locale.MonitoringLangfuseFormDescription
}
func (m *LangfuseFormModel) GetFormName() string {
return locale.MonitoringLangfuseFormName
}
func (m *LangfuseFormModel) GetFormSummary() string {
return ""
}
func (m *LangfuseFormModel) GetFormOverview() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(locale.MonitoringLangfuseFormTitle))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Bold(true).Render(locale.MonitoringLangfuseFormDescription))
sections = append(sections, "")
sections = append(sections, m.GetStyles().Paragraph.Render(locale.MonitoringLangfuseFormOverview))
return strings.Join(sections, "\n")
}
func (m *LangfuseFormModel) GetCurrentConfiguration() string {
var sections []string
sections = append(sections, m.GetStyles().Subtitle.Render(m.GetFormName()))
config := m.GetController().GetLangfuseConfig()
getMaskedValue := func(value string) string {
maskedValue := strings.Repeat("*", len(value))
if len(value) > 15 {
maskedValue = maskedValue[:15] + "..."
}
return maskedValue
}
switch config.DeploymentType {
case "embedded":
sections = append(sections, "• "+locale.UIMode+m.GetStyles().Success.Render(locale.MonitoringLangfuseEmbedded))
if listenIP := config.ListenIP.Value; listenIP != "" {
listenIP = m.GetStyles().Info.Render(listenIP)
sections = append(sections, fmt.Sprintf("• %s: %s", locale.MonitoringLangfuseListenIP, listenIP))
} else if listenIP := config.ListenIP.Default; listenIP != "" {
listenIP = m.GetStyles().Muted.Render(listenIP)
sections = append(sections, fmt.Sprintf("• %s: %s", locale.MonitoringLangfuseListenIP, listenIP))
}
if listenPort := config.ListenPort.Value; listenPort != "" {
listenPort = m.GetStyles().Info.Render(listenPort)
sections = append(sections, fmt.Sprintf("• %s: %s", locale.MonitoringLangfuseListenPort, listenPort))
} else if listenPort := config.ListenPort.Default; listenPort != "" {
listenPort = m.GetStyles().Muted.Render(listenPort)
sections = append(sections, fmt.Sprintf("• %s: %s", locale.MonitoringLangfuseListenPort, listenPort))
}
if config.BaseURL.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseBaseURL, m.GetStyles().Info.Render(config.BaseURL.Value)))
}
if config.ProjectID.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseProjectID, m.GetStyles().Info.Render(config.ProjectID.Value)))
}
if publicKey := config.PublicKey.Value; publicKey != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfusePublicKey, m.GetStyles().Muted.Render(getMaskedValue(publicKey))))
}
if secretKey := config.SecretKey.Value; secretKey != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseSecretKey, m.GetStyles().Muted.Render(getMaskedValue(secretKey))))
}
if config.AdminEmail.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseAdminEmail, m.GetStyles().Info.Render(config.AdminEmail.Value)))
}
if adminPassword := config.AdminPassword.Value; adminPassword != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseAdminPassword, m.GetStyles().Muted.Render(getMaskedValue(adminPassword))))
}
if config.AdminName.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseAdminName, m.GetStyles().Info.Render(config.AdminName.Value)))
}
case "external":
sections = append(sections, "• "+locale.UIMode+m.GetStyles().Success.Render(locale.MonitoringLangfuseExternal))
if config.BaseURL.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseBaseURL, m.GetStyles().Info.Render(config.BaseURL.Value)))
}
if config.ProjectID.Value != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseProjectID, m.GetStyles().Info.Render(config.ProjectID.Value)))
}
if publicKey := config.PublicKey.Value; publicKey != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfusePublicKey, m.GetStyles().Muted.Render(getMaskedValue(publicKey))))
}
if secretKey := config.SecretKey.Value; secretKey != "" {
sections = append(sections, fmt.Sprintf("• %s: %s",
locale.MonitoringLangfuseSecretKey, m.GetStyles().Muted.Render(getMaskedValue(secretKey))))
}
case "disabled":
sections = append(sections, "• "+locale.UIMode+m.GetStyles().Warning.Render(locale.MonitoringLangfuseDisabled))
}
return strings.Join(sections, "\n")
}
func (m *LangfuseFormModel) IsConfigured() bool {
config := m.GetController().GetLangfuseConfig()
return config.DeploymentType != "disabled"
}
func (m *LangfuseFormModel) GetHelpContent() string {
var sections []string
deploymentType := m.getSelectedDeploymentType()
sections = append(sections, m.GetStyles().Subtitle.Render(locale.MonitoringLangfuseFormTitle))
sections = append(sections, "")
sections = append(sections, locale.MonitoringLangfuseModeGuide)
sections = append(sections, "")
switch deploymentType {
case "embedded":
sections = append(sections, locale.MonitoringLangfuseEmbeddedHelp)
case "external":
sections = append(sections, locale.MonitoringLangfuseExternalHelp)
case "disabled":
sections = append(sections, locale.MonitoringLangfuseDisabledHelp)
}
return strings.Join(sections, "\n")
}
func (m *LangfuseFormModel) HandleSave() error {
config := m.GetController().GetLangfuseConfig()
deploymentType := m.getSelectedDeploymentType()
fields := m.GetFormFields()
// create a working copy of the current config to modify
newConfig := &controller.LangfuseConfig{
DeploymentType: deploymentType,
// copy current EnvVar fields - they preserve metadata like Line, IsPresent, etc.
ListenIP: config.ListenIP,
ListenPort: config.ListenPort,
BaseURL: config.BaseURL,
ProjectID: config.ProjectID,
PublicKey: config.PublicKey,
SecretKey: config.SecretKey,
AdminEmail: config.AdminEmail,
AdminPassword: config.AdminPassword,
AdminName: config.AdminName,
Installed: config.Installed,
LicenseKey: config.LicenseKey,
}
// update field values based on form input
for _, field := range fields {
value := strings.TrimSpace(field.Input.Value())
switch field.Key {
case "listen_ip":
newConfig.ListenIP.Value = value
case "listen_port":
newConfig.ListenPort.Value = value
case "base_url":
newConfig.BaseURL.Value = value
case "project_id":
newConfig.ProjectID.Value = value
case "public_key":
newConfig.PublicKey.Value = value
case "secret_key":
newConfig.SecretKey.Value = value
case "admin_email":
newConfig.AdminEmail.Value = value
case "admin_password":
newConfig.AdminPassword.Value = value
case "admin_name":
newConfig.AdminName.Value = value
case "license_key":
newConfig.LicenseKey.Value = value
}
}
// save the configuration
if err := m.GetController().UpdateLangfuseConfig(newConfig); err != nil {
logger.Errorf("[LangfuseFormModel] SAVE: error updating langfuse config: %v", err)
return err
}
logger.Log("[LangfuseFormModel] SAVE: success")
return nil
}
func (m *LangfuseFormModel) HandleReset() {
// reset config to defaults
config := m.GetController().ResetLangfuseConfig()
// reset deployment selection
m.GetListHelper().SelectByValue(&m.deploymentList, config.DeploymentType)
// rebuild form with reset deployment type
m.BuildForm()
}
func (m *LangfuseFormModel) OnFieldChanged(fieldIndex int, oldValue, newValue string) {
// additional validation could be added here if needed
}
func (m *LangfuseFormModel) GetFormFields() []FormField {
return m.BaseScreen.fields
}
func (m *LangfuseFormModel) SetFormFields(fields []FormField) {
m.BaseScreen.fields = fields
}
// BaseListHandler interface implementation
func (m *LangfuseFormModel) GetList() *list.Model {
return &m.deploymentList
}
func (m *LangfuseFormModel) GetListDelegate() *BaseListDelegate {
return m.deploymentDelegate
}
func (m *LangfuseFormModel) OnListSelectionChanged(oldSelection, newSelection string) {
// rebuild form when deployment type changes
m.BuildForm()
}
func (m *LangfuseFormModel) GetListTitle() string {
return locale.MonitoringLangfuseDeploymentType
}
func (m *LangfuseFormModel) GetListDescription() string {
return locale.MonitoringLangfuseDeploymentTypeDesc
}
// Update method - handle screen-specific input
func (m *LangfuseFormModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.KeyMsg:
// handle list input first (if focused on list)
if cmd := m.HandleListInput(msg); cmd != nil {
return m, cmd
}
// then handle field input
if cmd := m.HandleFieldInput(msg); cmd != nil {
return m, cmd
}
}
// delegate to base screen for common handling
cmd := m.BaseScreen.Update(msg)
return m, cmd
}
// Compile-time interface validation
var _ BaseScreenModel = (*LangfuseFormModel)(nil)
var _ BaseScreenHandler = (*LangfuseFormModel)(nil)
var _ BaseListHandler = (*LangfuseFormModel)(nil)
backend/cmd/installer/wizard/models/list_screen.go
+369
View File
@@ -0,0 +1,369 @@
package models
import (
"fmt"
"strings"
"pentagi/cmd/installer/wizard/controller"
"pentagi/cmd/installer/wizard/locale"
"pentagi/cmd/installer/wizard/styles"
"pentagi/cmd/installer/wizard/window"
"github.com/charmbracelet/bubbles/viewport"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
// ListScreenHandler defines methods that concrete list screens must implement
type ListScreenHandler interface {
// LoadItems loads the list items
LoadItems() []ListItem
// HandleSelection handles item selection, returns navigation command
HandleSelection(item ListItem) tea.Cmd
// GetOverview returns general overview content
GetOverview() string
// ShowConfiguredStatus returns whether to show configuration status icons
ShowConfiguredStatus() bool
}
// ListItem represents a single item in the list
type ListItem struct {
ID ScreenID
Model BaseScreenModel
Highlighted bool // if true, the item is the most important action for user
}
// ListScreen provides common functionality for menu/list screens
type ListScreen struct {
// Dependencies
controller controller.Controller
styles styles.Styles
window window.Window
registry Registry
// State
selectedIndex int
items []ListItem
// Handler
handler ListScreenHandler
}
// NewListScreen creates a new list screen instance
func NewListScreen(
c controller.Controller, s styles.Styles, w window.Window, r Registry, h ListScreenHandler,
) *ListScreen {
return &ListScreen{
controller: c,
styles: s,
window: w,
registry: r,
handler: h,
}
}
// BaseScreenModel interface partial implementation
func (l *ListScreen) GetFormOverview() string {
var sections []string
// general overview
if overview := l.handler.GetOverview(); overview != "" {
sections = append(sections, overview)
sections = append(sections, "")
}
// statistics
if len(l.items) > 0 {
if l.handler.ShowConfiguredStatus() {
configuredCount := 0
for _, item := range l.items {
if item.Model.IsConfigured() {
configuredCount++
}
}
sections = append(sections, l.styles.Subtitle.Render(locale.UIStatistics))
configuredText := l.styles.Success.Render(fmt.Sprintf("%s: %d", locale.StatusConfigured, configuredCount))
sections = append(sections, "• "+configuredText)
notConfiguredCount := len(l.items) - configuredCount
notConfiguredText := l.styles.Warning.Render(fmt.Sprintf("%s: %d", locale.StatusNotConfigured, notConfiguredCount))
sections = append(sections, "• "+notConfiguredText)
sections = append(sections, "")
}
}
return strings.Join(sections, "\n")
}
func (l *ListScreen) GetCurrentConfiguration() string {
var sections []string
for idx, item := range l.items {
sections = append(sections, item.Model.GetCurrentConfiguration())
if idx < len(l.items)-1 {
sections = append(sections, "")
}
}
return strings.Join(sections, "\n")
}
func (l *ListScreen) IsConfigured() bool {
var configuredCount int
for _, item := range l.items {
if item.Model.IsConfigured() {
configuredCount++
}
}
return configuredCount == len(l.items)
}
func (l *ListScreen) GetFormHotKeys() []string {
if len(l.items) > 0 {
return []string{"up|down", "enter"}
}
return []string{"enter"}
}
// tea.Model interface implementation
func (l *ListScreen) Init() tea.Cmd {
l.items = l.handler.LoadItems()
for i, item := range l.items {
if item.Model == nil {
l.items[i].Model = l.registry.GetScreen(item.ID)
}
}
if l.selectedIndex < 0 || l.selectedIndex >= len(l.items) {
l.selectedIndex = 0
}
return nil
}
func (l *ListScreen) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
// window resize handled by app.go
case tea.KeyMsg:
switch msg.String() {
case "up":
if l.selectedIndex > 0 {
l.selectedIndex--
}
case "down":
if l.selectedIndex < len(l.items)-1 {
l.selectedIndex++
}
case "enter":
return l.handleSelection()
}
}
return l, nil
}
func (l *ListScreen) View() string {
contentWidth, contentHeight := l.window.GetContentSize()
if contentWidth <= 0 || contentHeight <= 0 {
return locale.UILoading
}
leftPanel := l.renderItemsList()
rightPanel := l.renderItemInfo()
if l.isVerticalLayout() {
return l.renderVerticalLayout(leftPanel, rightPanel, contentWidth, contentHeight)
}
return l.renderHorizontalLayout(leftPanel, rightPanel, contentWidth, contentHeight)
}
// Helper methods for concrete implementations
// GetScreen returns the screen model for the given ID (implement Registry interface)
func (l *ListScreen) GetScreen(id ScreenID) BaseScreenModel {
return l.registry.GetScreen(id)
}
// GetController returns the state controller
func (l *ListScreen) GetController() controller.Controller {
return l.controller
}
// GetStyles returns the styles
func (l *ListScreen) GetStyles() styles.Styles {
return l.styles
}
// GetWindow returns the window
func (l *ListScreen) GetWindow() window.Window {
return l.window
}
// Internal methods
// handleSelection processes item selection
func (l *ListScreen) handleSelection() (tea.Model, tea.Cmd) {
if l.selectedIndex >= len(l.items) {
return l, nil
}
selectedItem := l.items[l.selectedIndex]
return l, l.handler.HandleSelection(selectedItem)
}
// getItemInfo is used as fallback only, most info comes from GetConfigScreen
func (l *ListScreen) getItemInfo(item ListItem) string {
var sections []string
sections = append(sections, l.styles.Subtitle.Render(item.Model.GetFormName()))
sections = append(sections, "")
sections = append(sections, l.styles.Paragraph.Render(item.Model.GetFormDescription()))
return strings.Join(sections, "\n")
}
// renderItemsList creates the left panel with items list
func (l *ListScreen) renderItemsList() string {
var sections []string
for i, item := range l.items {
selected := i == l.selectedIndex
var itemText string
if item.Model != nil {
if l.handler.ShowConfiguredStatus() {
statusIcon := l.styles.RenderStatusIcon(item.Model.IsConfigured()) + " "
itemText = statusIcon + item.Model.GetFormName()
} else {
itemText = item.Model.GetFormName()
}
} else {
// fallback to registry to resolve model for label
model := l.registry.GetScreen(item.ID)
if l.handler.ShowConfiguredStatus() {
statusIcon := l.styles.RenderStatusIcon(model.IsConfigured()) + " "
itemText = statusIcon + model.GetFormName()
} else {
itemText = model.GetFormName()
}
}
rendered := l.styles.RenderMenuItem(itemText, selected, false, item.Highlighted)
sections = append(sections, rendered)
}
if l.handler.ShowConfiguredStatus() {
sections = append(sections, "")
sections = append(sections, l.styles.Muted.Render(locale.LegendConfigured))
sections = append(sections, l.styles.Muted.Render(locale.LegendNotConfigured))
}
return strings.Join(sections, "\n")
}
// renderItemInfo creates the right panel with item details
func (l *ListScreen) renderItemInfo() string {
if len(l.items) == 0 || l.selectedIndex >= len(l.items) {
return l.styles.Info.Render(locale.UINoConfigSelected)
}
selectedItem := l.items[l.selectedIndex]
// try to get config screen overview first
if overview := selectedItem.Model.GetFormOverview(); overview != "" {
currentConfiguration := selectedItem.Model.GetCurrentConfiguration()
if currentConfiguration == "" {
return overview
}
wholeContent := overview + "\n" + currentConfiguration
if l.getContentTrueHeight(wholeContent)+PaddingHeight < l.window.GetContentHeight() {
return wholeContent
}
return overview
}
// fallback to handler's item info
return l.getItemInfo(selectedItem)
}
// Layout methods
func (l *ListScreen) getContentTrueHeight(content string) int {
contentWidth := l.window.GetContentWidth()
if l.isVerticalLayout() {
verticalStyle := lipgloss.NewStyle().Width(contentWidth).Padding(verticalLayoutPaddings...)
contentStyled := verticalStyle.Render(content)
return lipgloss.Height(contentStyled)
}
leftWidth, rightWidth := MinMenuWidth, MinInfoWidth
extraWidth := contentWidth - leftWidth - rightWidth - PaddingWidth
if extraWidth > 0 {
leftWidth = min(leftWidth+extraWidth/2, MaxMenuWidth)
rightWidth = contentWidth - leftWidth - PaddingWidth/2
}
contentStyled := lipgloss.NewStyle().Width(rightWidth).PaddingLeft(2).Render(content)
return lipgloss.Height(contentStyled)
}
// isVerticalLayout determines if vertical layout should be used
func (l *ListScreen) isVerticalLayout() bool {
contentWidth := l.window.GetContentWidth()
return contentWidth < (MinMenuWidth + MinInfoWidth + PaddingWidth)
}
// renderVerticalLayout renders content in vertical layout
func (l *ListScreen) renderVerticalLayout(leftPanel, rightPanel string, width, height int) string {
verticalStyle := lipgloss.NewStyle().Width(width).Padding(verticalLayoutPaddings...)
leftStyled := verticalStyle.Render(leftPanel)
rightStyled := verticalStyle.Render(rightPanel)
if lipgloss.Height(leftStyled)+lipgloss.Height(rightStyled)+3 < height {
return lipgloss.JoinVertical(lipgloss.Left,
verticalStyle.Render(leftPanel),
verticalStyle.Height(2).Render("\n"),
verticalStyle.Render(rightPanel),
)
}
return verticalStyle.Render(leftPanel)
}
// renderHorizontalLayout renders content in horizontal layout
func (l *ListScreen) renderHorizontalLayout(leftPanel, rightPanel string, width, height int) string {
leftWidth, rightWidth := MinMenuWidth, MinInfoWidth
extraWidth := width - leftWidth - rightWidth - PaddingWidth
if extraWidth > 0 {
leftWidth = min(leftWidth+extraWidth/2, MaxMenuWidth)
rightWidth = width - leftWidth - PaddingWidth/2
}
leftStyled := lipgloss.NewStyle().Width(leftWidth).Padding(horizontalLayoutPaddings...).Render(leftPanel)
rightStyled := lipgloss.NewStyle().Width(rightWidth).PaddingLeft(2).Render(rightPanel)
viewport := viewport.New(width, height-PaddingHeight)
viewport.SetContent(lipgloss.JoinHorizontal(lipgloss.Top, leftStyled, rightStyled))
return viewport.View()
}

Some files were not shown because too many files have changed in this diff Show More