Fix remote-host automation ownership (#6636)

Co-authored-by: Orca <help@stably.ai>

src/renderer/src/components/automations/AutomationsPage.tsx

@@ -167,6 +167,16 @@ function getAutomationHostTargetKey(target: AutomationHostTarget): string {
   return target.kind === 'environment' ? `environment:${target.environmentId}` : 'local'
 }
 
+function getAutomationHostTargetFromKey(key: string | null): AutomationHostTarget | null {
+  if (!key) {
+    return null
+  }
+  if (key.startsWith('environment:')) {
+    return { kind: 'environment', environmentId: key.slice('environment:'.length) }
+  }
+  return { kind: 'local' }
+}
+
 function getDefaultWorktree(worktrees: readonly Worktree[]): Worktree | null {
   return worktrees.find((worktree) => worktree.isMainWorktree) ?? worktrees[0] ?? null
 }
@@ -547,6 +557,10 @@ export default function AutomationsPage(): React.JSX.Element {
     () => worktreesByRepo[draft.projectId] ?? [],
     [draft.projectId, worktreesByRepo]
   )
+  const automationHostTarget = useMemo(
+    () => getAutomationHostTargetFromKey(automationHostTargetKey),
+    [automationHostTargetKey]
+  )
 
   useEffect(() => {
     for (const [workspaceId, worktree] of worktreeMap) {
@@ -801,6 +815,7 @@ export default function AutomationsPage(): React.JSX.Element {
         projectHostSetups,
         sshConnectionStates,
         runtimeStatusByEnvironmentId,
+        automationHostTarget,
         sourceHostAvailability: automationSourceHostAvailabilityById.get(selected.id)
       })
     : null
@@ -969,7 +984,7 @@ export default function AutomationsPage(): React.JSX.Element {
       pendingAutomationRunNavigation.hostId
         ? getAutomationTargetFromHostId(pendingAutomationRunNavigation.hostId)
         : selected
-          ? getAutomationOwnerTarget(selected)
+          ? getAutomationOwnerTarget(selected, automationHostTarget)
           : getAutomationListTarget(settings)
     void listAutomationRunsForTarget(target, automationId).then((nextRuns) => {
       if (!cancelled) {
@@ -979,7 +994,7 @@ export default function AutomationsPage(): React.JSX.Element {
     return () => {
       cancelled = true
     }
-  }, [pendingAutomationRunNavigation, selected, selected?.id, runs, settings])
+  }, [automationHostTarget, pendingAutomationRunNavigation, selected, selected?.id, runs, settings])
 
   useEffect(() => {
     const onAutomationsChanged = (): void => {
@@ -1478,7 +1493,7 @@ export default function AutomationsPage(): React.JSX.Element {
       }
       const automation = editingAutomationId
         ? currentAutomation
-          ? await updateAutomationForTarget(currentAutomation, updates)
+          ? await updateAutomationForTarget(currentAutomation, updates, automationHostTarget)
           : await window.api.automations.update({
               id: editingAutomationId,
               updates
@@ -1539,12 +1554,16 @@ export default function AutomationsPage(): React.JSX.Element {
   }
 
   const toggleAutomation = async (automation: Automation): Promise<void> => {
-    await updateAutomationForTarget(automation, { enabled: !automation.enabled })
+    await updateAutomationForTarget(
+      automation,
+      { enabled: !automation.enabled },
+      automationHostTarget
+    )
     await refresh()
   }
 
   const deleteAutomation = async (automation: Automation): Promise<void> => {
-    await deleteAutomationForTarget(automation)
+    await deleteAutomationForTarget(automation, automationHostTarget)
     if (useAppStore.getState().selectedAutomationId === automation.id) {
       selectAutomationId(null)
     }
@@ -1616,13 +1635,14 @@ export default function AutomationsPage(): React.JSX.Element {
       projectHostSetups,
       sshConnectionStates,
       runtimeStatusByEnvironmentId,
+      automationHostTarget,
       sourceHostAvailability: automationSourceHostAvailabilityById.get(automation.id)
     })
     if (!availability.canRunNow) {
       toast.error(availability.message)
       return
     }
-    await runAutomationNowForTarget(automation)
+    await runAutomationNowForTarget(automation, automationHostTarget)
     useAppStore.getState().recordFeatureInteraction('automation-run')
     await hydratePersistedUIState()
     await refresh()
@@ -1640,7 +1660,7 @@ export default function AutomationsPage(): React.JSX.Element {
     rerunRunIdsInFlightRef.current.add(runId)
     setRerunRunIdsInFlight(new Set(rerunRunIdsInFlightRef.current))
     try {
-      await runAutomationNowForTarget(automation)
+      await runAutomationNowForTarget(automation, automationHostTarget)
       await hydratePersistedUIState()
       await refresh()
       toast.message(
@@ -2217,6 +2237,7 @@ export default function AutomationsPage(): React.JSX.Element {
                 projectHostSetups,
                 sshConnectionStates,
                 runtimeStatusByEnvironmentId,
+                automationHostTarget,
                 sourceHostAvailability: automationSourceHostAvailabilityById.get(automation.id)
               })
               const workspaceLabel =

src/renderer/src/components/automations/automation-host-client.test.ts

@@ -4,7 +4,8 @@ import {
   createAutomationForTarget,
   getAutomationListTarget,
   listAutomationsForTarget,
-  runAutomationNowForTarget
+  runAutomationNowForTarget,
+  updateAutomationForTarget
 } from './automation-host-client'
 import { callRuntimeRpc } from '@/runtime/runtime-rpc-client'
 
@@ -126,4 +127,41 @@ describe('automation host client', () => {
       { timeoutMs: 15_000 }
     )
   })
+
+  it('updates and manually runs SSH-host automations through the remote server that listed them', async () => {
+    const automation = makeAutomation({
+      runContext: {
+        kind: 'workspace-run',
+        projectId: 'github:stablyai/orca',
+        hostId: 'ssh:devbox',
+        projectHostSetupId: 'setup-devbox',
+        repoId: 'repo-1',
+        path: '/srv/orca'
+      }
+    })
+    const sourceTarget = { kind: 'environment' as const, environmentId: 'gpu' }
+    vi.mocked(callRuntimeRpc)
+      .mockResolvedValueOnce({ automation: { ...automation, name: 'Updated' } })
+      .mockResolvedValueOnce({ run: { id: 'run-1', automationId: automation.id } })
+
+    await updateAutomationForTarget(automation, { name: 'Updated' }, sourceTarget)
+    await runAutomationNowForTarget(automation, sourceTarget)
+
+    expect(mockApi.automations.update).not.toHaveBeenCalled()
+    expect(mockApi.automations.runNow).not.toHaveBeenCalled()
+    expect(callRuntimeRpc).toHaveBeenNthCalledWith(
+      1,
+      sourceTarget,
+      'automation.update',
+      { id: automation.id, updates: { name: 'Updated' } },
+      { timeoutMs: 15_000 }
+    )
+    expect(callRuntimeRpc).toHaveBeenNthCalledWith(
+      2,
+      sourceTarget,
+      'automation.runNow',
+      { id: automation.id },
+      { timeoutMs: 15_000 }
+    )
+  })
 })

src/renderer/src/components/automations/automation-host-client.ts

@@ -43,8 +43,12 @@ export function getAutomationListTarget(
 }
 
 export function getAutomationOwnerTarget(
-  automation: Pick<Automation, 'runContext'>
+  automation: Pick<Automation, 'runContext'>,
+  sourceTarget?: AutomationHostTarget | null
 ): AutomationHostTarget {
+  if (sourceTarget?.kind === 'environment') {
+    return sourceTarget
+  }
   return getAutomationTargetFromHostId(automation.runContext?.hostId)
 }
 
@@ -121,9 +125,10 @@ export async function createAutomationForTarget(input: AutomationCreateInput): P
 
 export async function updateAutomationForTarget(
   automation: Automation,
-  updates: AutomationUpdateInput
+  updates: AutomationUpdateInput,
+  sourceTarget?: AutomationHostTarget | null
 ): Promise<Automation> {
-  const target = getAutomationOwnerTarget(automation)
+  const target = getAutomationOwnerTarget(automation, sourceTarget)
   if (target.kind === 'local') {
     return await window.api.automations.update({ id: automation.id, updates })
   }
@@ -136,8 +141,11 @@ export async function updateAutomationForTarget(
   return result.automation
 }
 
-export async function deleteAutomationForTarget(automation: Automation): Promise<void> {
-  const target = getAutomationOwnerTarget(automation)
+export async function deleteAutomationForTarget(
+  automation: Automation,
+  sourceTarget?: AutomationHostTarget | null
+): Promise<void> {
+  const target = getAutomationOwnerTarget(automation, sourceTarget)
   if (target.kind === 'local') {
     await window.api.automations.delete({ id: automation.id })
     return
@@ -145,8 +153,11 @@ export async function deleteAutomationForTarget(automation: Automation): Promise
   await callRuntimeRpc(target, 'automation.delete', { id: automation.id }, { timeoutMs: 15_000 })
 }
 
-export async function runAutomationNowForTarget(automation: Automation): Promise<AutomationRun> {
-  const target = getAutomationOwnerTarget(automation)
+export async function runAutomationNowForTarget(
+  automation: Automation,
+  sourceTarget?: AutomationHostTarget | null
+): Promise<AutomationRun> {
+  const target = getAutomationOwnerTarget(automation, sourceTarget)
   if (target.kind === 'local') {
     return await window.api.automations.runNow({ id: automation.id })
   }

src/renderer/src/components/automations/automation-target-availability.test.ts

@@ -140,6 +140,66 @@ describe('automation target availability', () => {
     ).toBe('host-mismatch')
   })
 
+  it('allows remote-listed SSH automations whose repo is projected through a runtime server', () => {
+    expect(
+      getAutomationTargetAvailability({
+        automation: makeAutomation({
+          executionTargetType: 'ssh',
+          executionTargetId: 'devbox',
+          runContext: {
+            kind: 'workspace-run',
+            projectId: 'project-1',
+            hostId: 'ssh:devbox',
+            projectHostSetupId: 'setup-1',
+            repoId: 'repo-1',
+            path: '/repo'
+          }
+        }),
+        repo: makeRepo({
+          connectionId: 'devbox',
+          executionHostId: 'runtime:gpu'
+        }),
+        workspace: makeWorkspace(),
+        projectHostSetups: [
+          makeProjectHostSetup({
+            hostId: 'ssh:devbox',
+            connectionId: 'devbox',
+            executionHostId: 'ssh:devbox'
+          })
+        ],
+        sshConnectionStates: new Map([['devbox', { status: 'connected' }]]),
+        automationHostTarget: { kind: 'environment', environmentId: 'gpu' }
+      })
+    ).toEqual({ canRunNow: true, reason: 'available', message: null })
+  })
+
+  it('allows remote-listed server-local automations whose setup is projected through a runtime server', () => {
+    expect(
+      getAutomationTargetAvailability({
+        automation: makeAutomation({
+          runContext: {
+            kind: 'workspace-run',
+            projectId: 'project-1',
+            hostId: 'local',
+            projectHostSetupId: 'setup-1',
+            repoId: 'repo-1',
+            path: '/repo'
+          }
+        }),
+        repo: makeRepo({ executionHostId: 'runtime:gpu' }),
+        workspace: makeWorkspace(),
+        projectHostSetups: [
+          makeProjectHostSetup({
+            hostId: 'runtime:gpu',
+            executionHostId: 'runtime:gpu'
+          })
+        ],
+        sshConnectionStates: new Map(),
+        automationHostTarget: { kind: 'environment', environmentId: 'gpu' }
+      })
+    ).toEqual({ canRunNow: true, reason: 'available', message: null })
+  })
+
   it('blocks saved run contexts whose project host setup is missing or not ready', () => {
     const automation = makeAutomation({
       runContext: {

src/renderer/src/components/automations/automation-target-availability.ts

@@ -8,6 +8,7 @@ import {
   MIN_COMPATIBLE_RUNTIME_SERVER_VERSION,
   RUNTIME_PROTOCOL_VERSION
 } from '../../../../shared/protocol-version'
+import type { AutomationHostTarget } from './automation-host-client'
 import type { SshConnectionState } from '../../../../shared/ssh-types'
 import type { TaskSourceContext } from '../../../../shared/task-source-context'
 import type { RuntimeStatus } from '../../../../shared/runtime-types'
@@ -52,6 +53,7 @@ type AutomationTargetAvailabilityArgs = {
     string,
     { status: RuntimeStatus | null; checkedAt: number }
   >
+  automationHostTarget?: AutomationHostTarget | null
   sourceHostAvailability?: readonly TaskSourceHostAvailability[]
 }
 
@@ -62,6 +64,7 @@ export function getAutomationTargetAvailability({
   projectHostSetups,
   sshConnectionStates,
   runtimeStatusByEnvironmentId,
+  automationHostTarget,
   sourceHostAvailability
 }: AutomationTargetAvailabilityArgs): AutomationTargetAvailability {
   if (!repo) {
@@ -93,15 +96,16 @@ export function getAutomationTargetAvailability({
         `Project setup on the selected automation host is ${setup.setupState}.`
       )
     }
-    if (
-      setup.projectId !== automation.runContext.projectId ||
-      setup.hostId !== automation.runContext.hostId ||
-      setup.repoId !== automation.runContext.repoId ||
-      setup.path !== automation.runContext.path ||
-      automation.runContext.repoId !== repo.id ||
-      automation.runContext.path !== repo.path ||
-      automation.runContext.hostId !== getRepoExecutionHostId(repo)
-    ) {
+    const setupMatchesContext =
+      setup.projectId === automation.runContext.projectId &&
+      setup.repoId === automation.runContext.repoId &&
+      setup.path === automation.runContext.path &&
+      setupHostMatchesRunContext(setup.hostId, automation.runContext.hostId, automationHostTarget)
+    const repoMatchesContext =
+      automation.runContext.repoId === repo.id &&
+      automation.runContext.path === repo.path &&
+      repoHostMatchesRunContext(repo, automation.runContext.hostId, automationHostTarget)
+    if (!setupMatchesContext || !repoMatchesContext) {
       return unavailable(
         'host-mismatch',
         'The saved run host no longer matches this project setup.'
@@ -142,6 +146,40 @@ export function getAutomationTargetAvailability({
   }
 }
 
+function getRuntimeTargetHostId(target: AutomationHostTarget | null | undefined): string | null {
+  return target?.kind === 'environment'
+    ? `runtime:${encodeURIComponent(target.environmentId)}`
+    : null
+}
+
+function setupHostMatchesRunContext(
+  setupHostId: string,
+  runHostId: string,
+  target: AutomationHostTarget | null | undefined
+): boolean {
+  if (setupHostId === runHostId) {
+    return true
+  }
+  const targetHostId = getRuntimeTargetHostId(target)
+  // Why: remote-runtime project lists project the server-local host as runtime:<env>,
+  // while CLI-created automations can preserve the server's durable local run host.
+  return targetHostId !== null && setupHostId === targetHostId && runHostId === 'local'
+}
+
+function repoHostMatchesRunContext(
+  repo: Repo,
+  runHostId: string,
+  target: AutomationHostTarget | null | undefined
+): boolean {
+  if (runHostId === getRepoExecutionHostId(repo)) {
+    return true
+  }
+  const targetHostId = getRuntimeTargetHostId(target)
+  // Why: repos fetched from a remote runtime are owned by runtime:<env> in the
+  // renderer, but saved automations still target the host setup that runs there.
+  return targetHostId !== null && getRepoExecutionHostId(repo) === targetHostId
+}
+
 function getAutomationSourceAvailability(
   sourceContext: TaskSourceContext | null | undefined,
   sourceHostAvailability: readonly TaskSourceHostAvailability[] | undefined