Desktop migration: Rust backend → Python backend (#5302)#5374
Desktop migration: Rust backend → Python backend (#5302)#5374
Conversation
Greptile SummaryThis PR migrates the Omi desktop backend from Rust to Python by adding new FastAPI endpoints and database helpers that parallel the existing Rust backend's capabilities. It introduces desktop-specific chat-session CRUD ( Key observation:
Confidence Score: 4/5
Sequence DiagramsequenceDiagram
participant Desktop as Desktop App (Swift)
participant API as FastAPI Backend (Python)
participant FS as Firestore
participant PC as Pinecone (background)
Note over Desktop,PC: Chat Session Flow
Desktop->>API: POST /v2/chat-sessions
API->>FS: add_chat_session(uid, session_data)
API-->>Desktop: ChatSessionResponse
Desktop->>API: POST /v2/messages/save
API->>FS: save_message(uid, message_data)
API->>FS: add_message_to_chat_session (preview, count, updated_at)
API-->>Desktop: SaveMessageResponse {id, created_at}
Desktop->>API: PATCH /v2/messages/{id}/rating
API->>FS: update_message_rating(uid, message_id, rating)
API-->>Desktop: StatusResponse {status: ok}
Desktop->>API: DELETE /v2/chat-sessions/{id}
API->>FS: delete_chat_session_messages (batch delete)
API->>FS: delete_chat_session
API-->>Desktop: StatusResponse {status: ok}
Note over Desktop,PC: Conversation from Segments Flow
Desktop->>API: POST /v1/conversations/from-segments
API->>API: validate segments, compute finished_at
API->>API: process_conversation(uid, language, obj)
API->>FS: store conversation
API-->>Desktop: FromSegmentsResponse {id, status, discarded}
Note over Desktop,PC: Screen Activity Sync Flow
Desktop->>API: POST /v1/screen-activity/sync
API->>FS: upsert_screen_activity (sync, truncates ocrText@1000)
API-->>Desktop: ScreenActivitySyncResponse {synced, last_id}
API-)PC: upsert_screen_activity_vectors (async daemon thread)
Last reviewed commit: 534a65f |
| class ChatSessionResponse(BaseModel): | ||
| id: str | ||
| title: str | ||
| preview: Optional[str] = None | ||
| created_at: datetime | ||
| updated_at: datetime | ||
| app_id: Optional[str] = None | ||
| message_count: int = 0 | ||
| starred: bool = False |
There was a problem hiding this comment.
ChatSessionResponse maps app_id directly, but existing Firestore chat-session documents only store plugin_id (the legacy field name). When get_chat_session_by_id or get_chat_sessions returns a document that has plugin_id but no app_id, Pydantic silently sets app_id=None in the response — even though the session does belong to an app.
New sessions written by create_chat_session (lines 578–579) correctly store both app_id and plugin_id, so they're fine. But any session created before this migration (or by the mobile app) will always report app_id=null to the desktop client.
A @model_validator can fix this:
| class ChatSessionResponse(BaseModel): | |
| id: str | |
| title: str | |
| preview: Optional[str] = None | |
| created_at: datetime | |
| updated_at: datetime | |
| app_id: Optional[str] = None | |
| message_count: int = 0 | |
| starred: bool = False | |
| from pydantic import BaseModel, model_validator | |
| class ChatSessionResponse(BaseModel): | |
| id: str | |
| title: str | |
| preview: Optional[str] = None | |
| created_at: datetime | |
| updated_at: datetime | |
| app_id: Optional[str] = None | |
| message_count: int = 0 | |
| starred: bool = False | |
| @model_validator(mode='before') | |
| @classmethod | |
| def _coerce_plugin_id(cls, data): | |
| if isinstance(data, dict) and data.get('app_id') is None: | |
| data['app_id'] = data.get('plugin_id') | |
| return data |
This ensures legacy sessions correctly surface their app association to the desktop client.
E2E Live Firestore Test — Day 2 Endpoints (21/21 PASS)Local backend on port 8789 running trunk Staged Tasks (8/8)Daily Scores (3/3)Focus Sessions (5/5)Advice (5/5)All 21 endpoints hit real Firestore and return expected responses. by AI for @beastoin |
E2E Live Test — Full Trunk (Day 1 + Day 2 + Auth)Local backend (23/23 PASS)Backend on port 8789 running trunk Mac Mini E2E (13/13 PASS)Python backend tunneled to Mac Mini ( All endpoints hit real Firestore. Integration PR ready for merge. by AI for @beastoin |
Day 3 Integration CompleteBoth sub-PRs merged into trunk:
Integration trunk status
Remaining before merge
by AI for @beastoin |
E2E Live Test Evidence (CP9) — kaiTested against local Python backend (localhost:8789) on A. New Desktop Endpoints (4/4 pass)
B. Adapted/Path-Changed Endpoints (4/4 pass)
C. Desktop Chat Session Flow (7/7 pass)
D. Mobile Regression Check (8/8 pass — no breakage)
Ren's Mac Mini Evidence (7/7 pass)OAuth sign-in OK, 7 endpoints verified via curl, desktop app launched successfully. All checkpoints passed (CP0-CP9). PR ready for merge. by AI for @beastoin |
Mac Mini Live Test — Python Backend Only (CP9)CONFIRMED: Desktop app runs fully on Python backend with zero Rust dependency. Evidence1. Authenticated app loading real data from Python backend: App shows beastoin's conversations loaded from dev Firestore via Python backend. Sidebar navigation, tasks, conversations all functional. 2. Backend logs — authenticated API calls from Mac Mini (100.126.187.125): All for UID R2IxlZVs8sRU20j9jLNTBiiFAoO2 (beastoin), all 200 OK. 3. OAuth flow verified through Python backend: 4. No Rust backend running:
5. Setup:
SummaryAll Day 1-3 endpoints work. Desktop app is fully operational on Python backend. Rust backend can be safely decommissioned. by AI for @beastoin |
f868760 to
f92eec0
Compare
Independent Verification — PR #5374Verifier: kelvin Test Results
Codex Audit
Cross-PR Interaction
Remote Sync
Verdict: PASS (after CRITICAL fixes applied) |
Independent Verification — PR #5374Verifier: noa (independent, did not author this code) Test Results
Codex Audit
Warnings (non-blocking)
Commands RunRemote Sync
Verdict: PASS |
Combined UAT Summary — Desktop Migration PRsVerifier: noa | Branch:
Combined: 1026 pass, 13 fail (pre-existing), 42 errors (env-only) | Cross-PR interference: none | Remote sync: verified Overall Verdict: PASS — ready for merge in order #5374 → #5395 → #5413 |
Replaces the hardcoded omi-desktop-auth Cloud Run URL with the OMI_API_URL environment variable, matching APIClient.baseURL resolution. Python backend already has identical /v1/auth/* endpoints. Closes #5359
Pass redirect_uri from the auth session to the callback HTML template instead of hardcoding omi://auth/callback. This enables desktop apps (which use omi-computer://auth/callback) to receive OAuth callbacks correctly when authenticating through the Python backend.
Both Google and Apple callback endpoints now pass the session's redirect_uri to the auth_callback.html template, enabling dynamic custom URL scheme redirects per client (mobile vs desktop).
Add server-side validation at /v1/auth/authorize to reject redirect_uri values that don't match allowed app schemes (omi://, omi-computer://, omi-computer-dev://). Also fix empty string fallback with 'or' operator.
Use |tojson filter for safe template variable serialization. Add defense-in-depth scheme validation in JavaScript before redirect. Block redirect and manual link for disallowed schemes.
…ering 15 tests covering: - Redirect_uri allowlist validation (rejects https, javascript, data, ftp, empty) - Allowed schemes pass (omi://, omi-computer://, omi-computer-dev://) - Google/Apple callback uses session redirect_uri in template - Fallback to default omi://auth/callback when missing - XSS safety: JSON-escaped redirect_uri prevents script injection
When FIREBASE_API_KEY has app restrictions (e.g. Android-only), the signInWithIdp REST API returns 403. Fall back to decoding the Google id_token JWT, looking up the user via Admin SDK (get_user_by_email), and creating a custom token directly. This makes auth work regardless of API key restrictions.
Desktop app uploads transcriptions via this endpoint but Python backend only had it in the developer API (API key auth). This adds a user-auth version to conversations router, reusing the same process_conversation pipeline. Defaults source to 'desktop', accepts timezone and input_device_name fields sent by Swift client.
New router for desktop app's session-based chat: - GET/POST/GET/:id/PATCH/:id/DELETE /v2/chat-sessions - POST /v2/desktop/messages (simple save, not streaming) - PATCH /v2/messages/:id/rating
…ions Path updates (5 endpoints): - v2/chat/initial-message → v2/initial-message - v2/agent/provision → v1/agent/vm-ensure - v2/agent/status → v1/agent/vm-status - v1/personas/check-username → v1/apps/check-username - v1/personas/generate-prompt → v1/app/generate-prompts (POST→GET) Decoder hardening: - ServerConversation.createdAt: use decodeIfPresent with Date() fallback - ActionItemsListResponse: try "action_items" then "items" key (Python vs staged-tasks) - AgentProvisionResponse/AgentStatusResponse: make fields optional, add hasVm - UsernameAvailableResponse: support both is_taken (Python) and available (Rust) Graceful no-ops: - recordLlmUsage(): no-op with log (endpoint removed) - fetchTotalOmiAICost(): return nil immediately (endpoint removed) - getChatMessageCount(): return 0 immediately (endpoint removed) Remove staged-tasks migration: - Remove migrateStagedTasks() and migrateConversationItemsToStaged() from APIClient - Remove migration callers and functions from TasksStore Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
C1: Replace unsafe base64 JWT decode with firebase_admin.auth.verify_id_token() which verifies signature against Google public keys before trusting claims. C2: Wrap email in sanitize_pii() per CLAUDE.md logging rules. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
94c9130 to
78d15d2
Compare
Independent Verification — PR #5374 (rebased)Verifier: noa | Branch: Test Results
Architecture Review
Mac Mini E2E (agent-swift + cliclick)
Warnings (non-blocking)
Verdict: ✅ PASS0 CRITICAL, 2 WARNING (non-blocking). Merge order: #5374 → #5395 → #5413. |
Deployment Steps ChecklistDeploy surfaces: Backend (Cloud Run + GKE backend-listen + GKE pusher) + Desktop (auto-deploy) Pre-merge
Backend deploy (hand to @mon)
Desktop deploy (automatic)
Post-deploy verification
Rollback plan
by AI for @beastoin |
Independent Verification — PR #5374 (collab/5302-integration)Verifier: noa (independent) Results
Notes
Verdict: PASS |
Independent Verification — PR #5374Verifier: noa (independent) ScopeDesktop auth + base integration: Google OAuth flow, AuthService, AppState, sidebar navigation, backend endpoints (auth, chat, advice, tasks, focus sessions, screen activity). Results
Codex Warnings (non-blocking)
Known Errors (pre-existing, not regressions)
Verdict: PASSZero regressions. Auth flow verified end-to-end with real Google OAuth. Warnings are non-blocking with clear follow-up paths. |
|
Agent VM gaps fixed and verified:
Test results:
All checkpoints passed (CP0-CP8). PR ready for merge pending manager approval. by AI for @beastoin |
Re-verification — PR #5374 (VM endpoint fix)Verifier: noa (independent) New Commits (8)
Code Review
Results
Verdict: PASSVM endpoints now complete — creation, status with restart, full response fields. All tests pass. |
Independent E2E Verification — Local BackendVerifier: noa (independent) Local Backend E2E TestSet up local Python backend from combined branch on Mac Mini, wired desktop app to Pipeline verified end-to-end: Results:
Evidence from app logs: Declarative E2E Flows (4/4 PASS)
Verdict: PASS — Desktop migration from Rust to Python backend verified end-to-end with local backend running combined PR code. Note: Current PR HEAD is 40ae983 — unit tests verified at that SHA in previous verification round. This E2E test was on 94c9130 (auth fix commit). |

Closes #5302. Migrates desktop from Rust backend to Python backend. Desktop becomes thin client; Python backend = source of truth.
What changed
Backend (kai) — 10 new REST endpoints: OAuth auth (Google/Apple + callback), chat sessions/messages + generate-title, conversations count, focus sessions CRUD, staged tasks CRUD, screen activity sync, advice CRUD, assistant settings. All use existing Firebase auth middleware.
Agent VM endpoints (kai) — Fixed 2 gaps found via Codex audit:
vm-ensurenow creates new GCE VMs for first-time users (ported from Rustagent.rs): generatesomi-agent-{uid[:12]}name, writes Firestore, spawns background GCE creation (e2-small, pd-ssd 50GB, startup script from GCS bucket).vm-statusnow returns all fields desktop needs (vm_name,ip,auth_token,zone,created_at,last_query_at) and triggers restart for stopped/terminated VMs (Rust parity).Desktop Swift (ren) — APIClient path fixes + decoder hardening, AuthService pointed to Python OAuth, TasksStore migration cleanup, TranscriptionRetryService path updates.
Tests — 147 PR-specific tests (134 original + 13 new agent VM tests covering vm-ensure creation, idempotency, restart, vm-status full fields, restart-on-check, GCE failure graceful degradation, UID truncation).
Architecture
Verification
78d15d27Agent VM tests: 13/13 PASS — vm-ensure creates new VMs, idempotent provisioning, restart stopped VMs, error recovery; vm-status returns full fields, triggers restart (Rust parity), graceful GCE failure; UID truncation to 12 chars.
Driver verdict: PASS. All endpoints tested live against dev Firestore on Mac Mini. No regressions.
Security
verify_id_token()(cryptographic JWT verification), not base64 decodesanitize_pii()on all email logging|tojsonfilter for JS contextuuid4(cryptographic randomness)Infra Prerequisites
https://api.omi.me/v1/auth/callback/googleand/apple) — same paths as existing Rust backendGOOGLE_CLIENT_ID,GOOGLE_CLIENT_SECRET,APPLE_CLIENT_ID,APPLE_TEAM_ID,APPLE_KEY_ID,APPLE_PRIVATE_KEY,BASE_API_URL,FIREBASE_API_KEYbased-hardware):GCE_PROJECT_ID,GCE_SOURCE_IMAGE,AGENT_GCS_BUCKETDeployment Steps
gh workflow run gcp_backend.yml -f environment=prod -f branch=main(Cloud Run image)gh workflow run gke_backend_listen.yml -f environment=prod -f branch=main(Helm rollout)gh workflow run gke_pusher.yml -f environment=prod -f branch=main(pusher image)desktop_auto_release.yml→ Codemagic/v4/desktop/chat,/v1/conversations/count,/v1/advice,/v1/auth/authorize,/v1/agent/vm-ensure,/v1/agent/vm-status), no new 5xx, desktop connects to Python backend./scripts/rollback_release.sh <tag>Merge order
This PR merges first → then #5395 → then #5413.
by AI for @beastoin