vik 0.1.1

Vik is an issue-driven coding workflow automation tool.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

//! Per-stage launch task.
//!
//! Each call to [`StageLauncher::launch`] spawns an independent task that
//! takes a stage from `before_run` through session monitoring and
//! `after_run`. The launcher carries no running-state; every transition
//! is reported back through the event channel so the main loop is the
//! sole owner of [`super::running::RunningMap`].

use std::sync::Arc;

use thiserror::Error;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;

use crate::context::IssueStage;
use crate::hooks::HookError;
use crate::logging::stage_span;
use crate::session::{SessionError, SessionFactory, SessionState};
use crate::workflow::Workflow;

use super::event::EventProducer;
use super::monitor::SessionMonitor;

#[derive(Clone)]
pub(super) struct StageLauncher {
  workflow: Arc<Workflow>,
  factory: SessionFactory,
  producer: EventProducer,
}

impl StageLauncher {
  pub(super) fn new(workflow: Arc<Workflow>, factory: SessionFactory, producer: EventProducer) -> Self {
    Self {
      workflow,
      factory,
      producer,
    }
  }

  /// Spawn the launch task. The runtime `IssueStage` already carries
  /// the issue run context needed for hooks and session spawn.
  pub(super) fn launch(&self, issue_stage: IssueStage, shutdown: CancellationToken) {
    let span = stage_span(issue_stage.stage_name(), &issue_stage.stage().agent);
    let launcher = self.clone();

    tokio::spawn(async move { launcher.run(issue_stage, shutdown).await }.instrument(span));
  }

  async fn run(self, issue_stage: IssueStage, shutdown: CancellationToken) {
    if shutdown.is_cancelled() {
      return;
    }

    let key = issue_stage.key();
    let session = match self.start_session(&issue_stage).await {
      Ok(session) => session,
      Err(error) => {
        tracing::error!(error = %error, "Failed to start session for issue stage");
        // No session ever existed; report under the reserved key so the
        // main loop can release the reservation.
        self.producer.stage_failed(key, error).await;
        return;
      },
    };

    self.producer.stage_started(issue_stage.clone(), session.clone()).await;

    let monitor = SessionMonitor::new(key.clone(), session.clone(), self.producer.clone());
    let terminal = tokio::select! {
      snapshot = monitor.watch() => snapshot,
      _ = shutdown.cancelled() => {
        // Cooperative shutdown: cancel the child then wait for the
        // session's own state machine to mark itself Cancelled. Without
        // the wait, `terminal` could observe a stale state.
        session.cancel();
        session.wait().await
      }
    };

    // `after_run` is skipped on cancellation only — failures still want
    // their cleanup to run. The hook's own failure is logged but not
    // propagated, so a flaky `after_run` cannot mask a stage result.
    if !matches!(terminal.state, SessionState::Cancelled)
      && let Err(error) = self
        .workflow
        .hooks()
        .after_issue_stage_run(&issue_stage, &issue_stage.stage().hooks.after_run)
        .await
    {
      tracing::error!(
        error = %error,
        "issue stage after_run hook failed",
      );
    }

    self.producer.stage_terminal(key, terminal).await;
  }

  async fn start_session(&self, issue_stage: &IssueStage) -> Result<crate::session::Session, StageLaunchError> {
    // `before_run` failure aborts the stage before any agent process is
    // spawned, so a misconfigured setup hook cannot waste tokens.
    if let Err(err) = self
      .workflow
      .hooks()
      .before_issue_stage_run(issue_stage, &issue_stage.stage().hooks.before_run)
      .await
    {
      tracing::error!(
        error = %err,
        "issue stage before_run hook failed",
      );
      return Err(StageLaunchError::Hook(err));
    }

    Ok(self.factory.spawn_stage(issue_stage.clone()).await?)
  }
}

#[derive(Debug, Error)]
enum StageLaunchError {
  #[error(transparent)]
  Hook(#[from] HookError),
  #[error(transparent)]
  Session(#[from] SessionError),
}