// Copyright 2026 Alibaba Group // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Package edition provides an extension-point mechanism that allows private // overlay modules (e.g. an internal distribution) to customise CLI behaviour // without modifying the open-source core. The open-source build uses the // zero-value defaults; an overlay calls Override before Execute. package edition import ( "context" "sync" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/agentproduct" "github.com/spf13/cobra" ) // ServerInfo describes a static MCP server endpoint. type ServerInfo struct { ID string Name string Endpoint string Prefixes []string } // ContentBlock is a single content item in a ToolResult. type ContentBlock struct { Type string Text string } // ToolResult holds the response from an MCP tool call. type ToolResult struct { Content []ContentBlock } // ToolCaller abstracts MCP tool invocation so private overlays can call MCP // tools without importing internal packages. The open-source core provides a // concrete adapter wrapping executor.Runner. type ToolCaller interface { // CallTool invokes an MCP tool by product ID and tool name. CallTool(ctx context.Context, productID, toolName string, args map[string]any) (*ToolResult, error) // Format returns the current output format ("json", "table", "raw"). Format() string // DryRun returns true when --dry-run is active. DryRun() bool // Fields returns the global --fields output projection ("" if unset). Fields() string // JQ returns the global --jq output filter expression ("" if unset). JQ() string } // ReadToolCaller is an optional capability for a narrowly classified read // lookup that must still execute while the outer command is rendering a // dry-run plan. Implementations must fail closed unless they can bypass the // global write barrier without weakening it for ordinary CallTool calls. // // The Shortcut runtime uses this only after accepting a narrow read-only tool // name classification. // Keeping it separate from ToolCaller means existing callers remain protected // by the default "dry-run executes nothing" contract. type ReadToolCaller interface { CallReadTool(ctx context.Context, productID, toolName string, args map[string]any) (*ToolResult, error) } // RuntimeDefaultFn resolves a single runtimeDefault placeholder (e.g. // "$currentUserId") to a concrete string value. Called lazily at RunE time. // Returning (_, false) is equivalent to "not registered" and falls through // to the next-lower default source. type RuntimeDefaultFn func(ctx context.Context) (string, bool) // Hooks groups all edition-specific behavioural overrides. Zero values // fall back to open-source defaults so the struct is safe to use as-is. type Hooks struct { // --- identity --- Name string // "open" (default) / overlay identifier ScenarioCode string // injected into x-dingtalk-scenario-code header // ClawTypeValue is the display identity carried only in message-send tool // arguments (parameter clawType) so the IM server can render the // "Send from AI" indicator on delivered messages. A valid non-empty // DWS_AGENT_PRODUCT overrides this display default, but never the separate // HTTP claw-type routing/PAT header. Empty → DefaultOSSClawType; overlays // set their own message-display default (e.g. "wukong"). ClawTypeValue string // PersonalEventSourceID identifies the personal-event source channel // used by dws event --as user. Empty → "open". Overlays set their own // value through Override; open-source core never infers custom sources // from environment/user input. PersonalEventSourceID string // --- runtime mode --- IsEmbedded bool // true when running inside a host application HideAuthLogin bool // true suppresses the "dws auth login" command AutoPurgeToken bool // true deletes local token data on expiry // --- paths --- ConfigDir func() string // custom config directory; nil → ~/.dws // --- HTTP headers --- // MergeHeaders must preserve base headers. If it sets claw-type, that value // must be deterministic and independent of the supplied base map because // PAT error serialization resolves it with an empty map. The hook must not // perform network, keychain, credential-refresh, or other blocking work. MergeHeaders func(base map[string]string) map[string]string // --- EnterpriseCredential HTTP headers --- // This hook is only for credential material. It must not set claw-type or // x-dws-agent-product; the core reasserts both after the hook returns. EnterpriseCredentialHeaders func(base map[string]string) map[string]string // --- auth --- AuthClientID string // OAuth client ID for device-flow authorisation AuthClientFromMCP bool // true → fetch client ID from MCP at runtime OnAuthError func(configDir string, err error) error TokenProvider func(ctx context.Context, fallback func() (string, error)) (string, error) // --- token persistence (overlay-managed keychain / encrypted storage) --- SaveToken func(configDir string, data []byte) error // persist token blob LoadToken func(configDir string) ([]byte, error) // retrieve token blob DeleteToken func(configDir string) error // remove persisted token // --- MCP result classification --- // ClassifyToolResult inspects raw MCP tool-call content and returns a typed // error (e.g. PATError, CLIError) when the response contains a known // gateway-auth or PAT-permission failure. nil → no special handling. // ClassifyToolResult func(content map[string]any) error // --- product & endpoint --- StaticServers func() []ServerInfo // non-nil → skip Market discovery VisibleProducts func() []string // non-nil → override help visibility RegisterExtraCommands func(root *cobra.Command, caller ToolCaller) // register overlay-only commands // --- discovery --- // DiscoveryURL overrides the Market API endpoint for server list. // Non-empty → loadDynamicCommands uses FetchServersFromURL(DiscoveryURL) // instead of the default Market base URL. Provides edition-level isolation. DiscoveryURL string // DiscoveryHeaders returns HTTP headers injected into discovery requests. // Used to authenticate edition-specific endpoints. DiscoveryHeaders func() map[string]string // SupplementServers returns edition-specific MCP servers NOT registered // in any Market registry. Always merged into the endpoint map alongside // Market/cache results, regardless of discovery success or failure. SupplementServers func() []ServerInfo // FallbackServers returns the full server list as a safety net when // Market discovery + cache both fail. Results are NOT cached so the // next startup still attempts live discovery. FallbackServers func() []ServerInfo // AfterPersistentPreRun runs at the end of the root PersistentPreRunE after // global setup (OAuth flag overrides, log level, output sink). Overlays use // this for clients that bypass the MCP runner (e.g. A2A gateway). AfterPersistentPreRun func(cmd *cobra.Command, args []string) error // ClassifyToolResult is called before the framework's default business-error // detection on MCP tool results. If it returns a non-nil error, that error // is used instead of the generic CategoryAPI business error. Editions use // this to return custom error types with specific exit codes (e.g. PAT // authorization errors with exit code 4). ClassifyToolResult func(content map[string]any) error // --- schema v3: runtime defaults --- // RuntimeDefaults returns resolvers for runtimeDefault placeholders (e.g. // "$currentUserId" → fn). Placeholders not in the map fall through to a // "not registered" warning. Open-source core returns an empty map; // overlays populate the whitelist ($currentUserId / $unionId / $corpId / // $now / $today). See schema v3 §2.3. RuntimeDefaults func() map[string]RuntimeDefaultFn } var ( mu sync.RWMutex current = defaultHooks() ) // Get returns the active edition hooks (never nil). func Get() *Hooks { mu.RLock() defer mu.RUnlock() return current } // Override replaces the active edition hooks. Must be called before Execute. func Override(h *Hooks) { if h == nil { return } mu.Lock() defer mu.Unlock() current = h } // ClawType returns the message-display identity for the active edition. // A valid non-empty DWS_AGENT_PRODUCT wins; otherwise the active edition's // ClawTypeValue (or DefaultOSSClawType) is used. Message-send helpers attach // this value as the clawType tool argument so the IM server can label delivered // messages as sent via AI. It never changes the HTTP claw-type routing/PAT // header. func ClawType() string { fallback := Get().ClawTypeValue if fallback == "" { fallback = DefaultOSSClawType } value, err := agentproduct.ResolveFromEnv(fallback) if err != nil { return fallback } return value } // PersonalEventSourceID returns the source channel for user-level events. func PersonalEventSourceID() string { if v := Get().PersonalEventSourceID; v != "" { return v } return "open" }